disable some CategoricalIndex comparisons

Author: jreback

pandas/core/index.py

@@ -45,27 +45,6 @@ def _try_get_item(x):
     except AttributeError:
         return x
 
-def _indexOp(opname):
-    """
-    Wrapper function for index comparison operations, to avoid
-    code duplication.
-    """
-
-    def wrapper(self, other):
-        func = getattr(self._data.view(np.ndarray), opname)
-        result = func(np.asarray(other))
-
-        # technically we could support bool dtyped Index
-        # for now just return the indexing array directly
-        if is_bool_dtype(result):
-            return result
-        try:
-            return Index(result)
-        except:  # pragma: no cover
-            return result
-    return wrapper
-
-
 class InvalidIndexError(Exception):
     pass
 
@@ -1216,13 +1195,6 @@ def __sub__(self, other):
                           "use .difference()",FutureWarning)
         return self.difference(other)
 
-    __eq__ = _indexOp('__eq__')
-    __ne__ = _indexOp('__ne__')
-    __lt__ = _indexOp('__lt__')
-    __gt__ = _indexOp('__gt__')
-    __le__ = _indexOp('__le__')
-    __ge__ = _indexOp('__ge__')
-
     def __and__(self, other):
         return self.intersection(other)
 
@@ -2380,6 +2352,34 @@ def _evaluate_with_timedelta_like(self, other, op, opstr):
     def _evaluate_with_datetime_like(self, other, op, opstr):
         raise TypeError("can only perform ops with datetime like values")
 
+    @classmethod
+    def _add_comparison_methods(cls):
+        """ add in comparison methods """
+
+        def _make_compare(op):
+
+            def _evaluate_compare(self, other):
+                func = getattr(self._data.view(np.ndarray), op)
+                result = func(np.asarray(other))
+
+                # technically we could support bool dtyped Index
+                # for now just return the indexing array directly
+                if is_bool_dtype(result):
+                    return result
+                try:
+                    return Index(result)
+                except:  # pragma: no cover
+                    return result
+
+            return _evaluate_compare
+
+        cls.__eq__ = _make_compare('__eq__')
+        cls.__ne__ = _make_compare('__ne__')
+        cls.__lt__ = _make_compare('__lt__')
+        cls.__gt__ = _make_compare('__gt__')
+        cls.__le__ = _make_compare('__le__')
+        cls.__ge__ = _make_compare('__ge__')
+
     @classmethod
     def _add_numeric_methods_disabled(cls):
         """ add in numeric methods to disable """
@@ -2530,6 +2530,7 @@ def invalid_op(self, other=None):
 
 Index._add_numeric_methods_disabled()
 Index._add_logical_methods()
+Index._add_comparison_methods()
 
 class CategoricalIndex(Index):
     """
@@ -2817,8 +2818,42 @@ def convert(c):
         cat = Categorical.from_codes(codes, categories=categories)
         return CategoricalIndex(cat, name=name)
 
+    @classmethod
+    def _add_comparison_methods(cls):
+        """ add in comparison methods """
+
+        def _make_compare(op):
+
+            def _evaluate_compare(self, other):
+
+                # we must have only CategoricalIndexes here
+                if not isinstance(other, CategoricalIndex):
+                    raise TypeError("cannot compare a non-categorical index vs a categorical index")
+                if not other.categories.equals(self.categories):
+                    raise TypeError("categorical index comparisions must have the same categories")
+
+                return getattr(self.codes, op)(other.codes)
+
+            return _evaluate_compare
+
+        def _make_invalid_op(name):
+
+            def invalid_op(self, other=None):
+                raise TypeError("cannot perform {name} with this index type: {typ}".format(name=name,
+                                                                                           typ=type(self)))
+            invalid_op.__name__ = name
+            return invalid_op
+
+        cls.__eq__ = _make_compare('__eq__')
+        cls.__ne__ = _make_compare('__ne__')
+        cls.__lt__ = _make_invalid_op('__lt__')
+        cls.__gt__ = _make_invalid_op('__gt__')
+        cls.__le__ = _make_invalid_op('__le__')
+        cls.__ge__ = _make_invalid_op('__ge__')
+
 CategoricalIndex._add_numeric_methods_disabled()
 CategoricalIndex._add_logical_methods_disabled()
+CategoricalIndex._add_comparison_methods()
 
 class NumericIndex(Index):
     """

pandas/tests/test_index.py

@@ -1484,6 +1484,29 @@ def test_get_indexer(self):
         self.assertRaises(NotImplementedError, lambda : idx2.get_indexer(idx1, method='backfill'))
         self.assertRaises(NotImplementedError, lambda : idx2.get_indexer(idx1, method='nearest'))
 
+    def test_equals(self):
+
+        ci1 = CategoricalIndex(['a', 'b'], categories=['a', 'b'])
+        ci2 = CategoricalIndex(['a', 'b'], categories=['a', 'b', 'c'])
+
+        self.assertTrue(ci1.equals(ci1))
+        self.assertFalse(ci1.equals(ci2))
+        self.assertTrue((ci1 == ci1).all())
+        self.assertFalse((ci1 != ci1).all())
+
+        # invalid comparisons
+        self.assertRaises(TypeError, lambda : ci1 == 1)
+        self.assertRaises(TypeError, lambda : ci1 == Index(['a','b']))
+        self.assertRaises(TypeError, lambda : ci1 == Index(['a','b']))
+        self.assertRaises(TypeError, lambda : ci1 == ci2)
+
+        # invalid ops
+        self.assertRaises(TypeError, lambda : ci1 < ci1)
+        self.assertRaises(TypeError, lambda : ci1 <= ci1)
+        self.assertRaises(TypeError, lambda : ci1 > ci1)
+        self.assertRaises(TypeError, lambda : ci1 >= ci1)
+
+
 class Numeric(Base):
 
     def test_numeric_compat(self):