ENH: enable mul, div on Index by dispatching to Series (#34160)

Author: jbrockmendel

doc/source/whatsnew/v1.2.0.rst

@@ -17,7 +17,7 @@ Enhancements
 
 Other enhancements
 ^^^^^^^^^^^^^^^^^^
-
+- :class:`Index` with object dtype supports division and multiplication (:issue:`34160`)
 -
 -
 

pandas/core/indexes/base.py

@@ -2377,31 +2377,10 @@ def _get_unique_index(self, dropna: bool = False):
     # --------------------------------------------------------------------
     # Arithmetic & Logical Methods
 
-    def __add__(self, other):
-        if isinstance(other, (ABCSeries, ABCDataFrame)):
-            return NotImplemented
-        from pandas import Series
-
-        return Index(Series(self) + other)
-
-    def __radd__(self, other):
-        from pandas import Series
-
-        return Index(other + Series(self))
-
     def __iadd__(self, other):
         # alias for __add__
         return self + other
 
-    def __sub__(self, other):
-        return Index(np.array(self) - other)
-
-    def __rsub__(self, other):
-        # wrap Series to ensure we pin name correctly
-        from pandas import Series
-
-        return Index(other - Series(self))
-
     def __and__(self, other):
         return self.intersection(other)
 
@@ -5293,38 +5272,6 @@ def _add_comparison_methods(cls):
         cls.__le__ = _make_comparison_op(operator.le, cls)
         cls.__ge__ = _make_comparison_op(operator.ge, cls)
 
-    @classmethod
-    def _add_numeric_methods_add_sub_disabled(cls):
-        """
-        Add in the numeric add/sub methods to disable.
-        """
-        cls.__add__ = make_invalid_op("__add__")
-        cls.__radd__ = make_invalid_op("__radd__")
-        cls.__iadd__ = make_invalid_op("__iadd__")
-        cls.__sub__ = make_invalid_op("__sub__")
-        cls.__rsub__ = make_invalid_op("__rsub__")
-        cls.__isub__ = make_invalid_op("__isub__")
-
-    @classmethod
-    def _add_numeric_methods_disabled(cls):
-        """
-        Add in numeric methods to disable other than add/sub.
-        """
-        cls.__pow__ = make_invalid_op("__pow__")
-        cls.__rpow__ = make_invalid_op("__rpow__")
-        cls.__mul__ = make_invalid_op("__mul__")
-        cls.__rmul__ = make_invalid_op("__rmul__")
-        cls.__floordiv__ = make_invalid_op("__floordiv__")
-        cls.__rfloordiv__ = make_invalid_op("__rfloordiv__")
-        cls.__truediv__ = make_invalid_op("__truediv__")
-        cls.__rtruediv__ = make_invalid_op("__rtruediv__")
-        cls.__mod__ = make_invalid_op("__mod__")
-        cls.__divmod__ = make_invalid_op("__divmod__")
-        cls.__neg__ = make_invalid_op("__neg__")
-        cls.__pos__ = make_invalid_op("__pos__")
-        cls.__abs__ = make_invalid_op("__abs__")
-        cls.__inv__ = make_invalid_op("__inv__")
-
     @classmethod
     def _add_numeric_methods_binary(cls):
         """
@@ -5340,11 +5287,12 @@ def _add_numeric_methods_binary(cls):
         cls.__truediv__ = _make_arithmetic_op(operator.truediv, cls)
         cls.__rtruediv__ = _make_arithmetic_op(ops.rtruediv, cls)
 
-        # TODO: rmod? rdivmod?
         cls.__mod__ = _make_arithmetic_op(operator.mod, cls)
+        cls.__rmod__ = _make_arithmetic_op(ops.rmod, cls)
         cls.__floordiv__ = _make_arithmetic_op(operator.floordiv, cls)
         cls.__rfloordiv__ = _make_arithmetic_op(ops.rfloordiv, cls)
         cls.__divmod__ = _make_arithmetic_op(divmod, cls)
+        cls.__rdivmod__ = _make_arithmetic_op(ops.rdivmod, cls)
         cls.__mul__ = _make_arithmetic_op(operator.mul, cls)
         cls.__rmul__ = _make_arithmetic_op(ops.rmul, cls)
 
@@ -5504,7 +5452,7 @@ def shape(self):
         return self._values.shape
 
 
-Index._add_numeric_methods_disabled()
+Index._add_numeric_methods()
 Index._add_logical_methods()
 Index._add_comparison_methods()
 

pandas/core/indexes/category.py

@@ -770,6 +770,4 @@ def _wrap_joined_index(
         return self._create_from_codes(joined, name=name)
 
 
-CategoricalIndex._add_numeric_methods_add_sub_disabled()
-CategoricalIndex._add_numeric_methods_disabled()
 CategoricalIndex._add_logical_methods_disabled()

pandas/core/indexes/datetimes.py

@@ -842,7 +842,6 @@ def indexer_between_time(
         return mask.nonzero()[0]
 
 
-DatetimeIndex._add_numeric_methods_disabled()
 DatetimeIndex._add_logical_methods_disabled()
 
 

pandas/core/indexes/multi.py

@@ -50,6 +50,7 @@
 from pandas.core.indexes.frozen import FrozenList
 from pandas.core.indexes.numeric import Int64Index
 import pandas.core.missing as missing
+from pandas.core.ops.invalid import make_invalid_op
 from pandas.core.sorting import (
     get_group_index,
     indexer_from_factorized,
@@ -3606,6 +3607,40 @@ def isin(self, values, level=None):
                 return np.zeros(len(levs), dtype=np.bool_)
             return levs.isin(values)
 
+    @classmethod
+    def _add_numeric_methods_add_sub_disabled(cls):
+        """
+        Add in the numeric add/sub methods to disable.
+        """
+        cls.__add__ = make_invalid_op("__add__")
+        cls.__radd__ = make_invalid_op("__radd__")
+        cls.__iadd__ = make_invalid_op("__iadd__")
+        cls.__sub__ = make_invalid_op("__sub__")
+        cls.__rsub__ = make_invalid_op("__rsub__")
+        cls.__isub__ = make_invalid_op("__isub__")
+
+    @classmethod
+    def _add_numeric_methods_disabled(cls):
+        """
+        Add in numeric methods to disable other than add/sub.
+        """
+        cls.__pow__ = make_invalid_op("__pow__")
+        cls.__rpow__ = make_invalid_op("__rpow__")
+        cls.__mul__ = make_invalid_op("__mul__")
+        cls.__rmul__ = make_invalid_op("__rmul__")
+        cls.__floordiv__ = make_invalid_op("__floordiv__")
+        cls.__rfloordiv__ = make_invalid_op("__rfloordiv__")
+        cls.__truediv__ = make_invalid_op("__truediv__")
+        cls.__rtruediv__ = make_invalid_op("__rtruediv__")
+        cls.__mod__ = make_invalid_op("__mod__")
+        cls.__rmod__ = make_invalid_op("__rmod__")
+        cls.__divmod__ = make_invalid_op("__divmod__")
+        cls.__rdivmod__ = make_invalid_op("__rdivmod__")
+        cls.__neg__ = make_invalid_op("__neg__")
+        cls.__pos__ = make_invalid_op("__pos__")
+        cls.__abs__ = make_invalid_op("__abs__")
+        cls.__inv__ = make_invalid_op("__inv__")
+
 
 MultiIndex._add_numeric_methods_disabled()
 MultiIndex._add_numeric_methods_add_sub_disabled()

pandas/core/indexes/period.py

@@ -724,7 +724,6 @@ def memory_usage(self, deep=False):
         return result
 
 
-PeriodIndex._add_numeric_methods_disabled()
 PeriodIndex._add_logical_methods_disabled()
 
 

pandas/tests/arithmetic/test_numeric.py

@@ -548,20 +548,6 @@ class TestMultiplicationDivision:
     # __mul__, __rmul__, __div__, __rdiv__, __floordiv__, __rfloordiv__
     # for non-timestamp/timedelta/period dtypes
 
-    @pytest.mark.parametrize(
-        "box",
-        [
-            pytest.param(
-                pd.Index,
-                marks=pytest.mark.xfail(
-                    reason="Index.__div__ always raises", raises=TypeError
-                ),
-            ),
-            pd.Series,
-            pd.DataFrame,
-        ],
-        ids=lambda x: x.__name__,
-    )
     def test_divide_decimal(self, box):
         # resolves issue GH#9787
         ser = Series([Decimal(10)])

pandas/tests/indexes/categorical/test_category.py

@@ -43,7 +43,14 @@ def test_disallow_addsub_ops(self, func, op_name):
         # GH 10039
         # set ops (+/-) raise TypeError
         idx = pd.Index(pd.Categorical(["a", "b"]))
-        msg = f"cannot perform {op_name} with this index type: CategoricalIndex"
+        cat_or_list = "'(Categorical|list)' and '(Categorical|list)'"
+        msg = "|".join(
+            [
+                f"cannot perform {op_name} with this index type: CategoricalIndex",
+                "can only concatenate list",
+                rf"unsupported operand type\(s\) for [\+-]: {cat_or_list}",
+            ]
+        )
         with pytest.raises(TypeError, match=msg):
             func(idx)
 

pandas/tests/indexes/common.py

@@ -146,22 +146,41 @@ def test_numeric_compat(self):
         # Check that this doesn't cover MultiIndex case, if/when it does,
         #  we can remove multi.test_compat.test_numeric_compat
         assert not isinstance(idx, MultiIndex)
+        if type(idx) is Index:
+            return
 
-        with pytest.raises(TypeError, match="cannot perform __mul__"):
+        typ = type(idx._data).__name__
+        lmsg = "|".join(
+            [
+                rf"unsupported operand type\(s\) for \*: '{typ}' and 'int'",
+                "cannot perform (__mul__|__truediv__|__floordiv__) with "
+                f"this index type: {typ}",
+            ]
+        )
+        with pytest.raises(TypeError, match=lmsg):
             idx * 1
-        with pytest.raises(TypeError, match="cannot perform __rmul__"):
+        rmsg = "|".join(
+            [
+                rf"unsupported operand type\(s\) for \*: 'int' and '{typ}'",
+                "cannot perform (__rmul__|__rtruediv__|__rfloordiv__) with "
+                f"this index type: {typ}",
+            ]
+        )
+        with pytest.raises(TypeError, match=rmsg):
             1 * idx
 
-        div_err = "cannot perform __truediv__"
+        div_err = lmsg.replace("*", "/")
         with pytest.raises(TypeError, match=div_err):
             idx / 1
-
-        div_err = div_err.replace(" __", " __r")
+        div_err = rmsg.replace("*", "/")
         with pytest.raises(TypeError, match=div_err):
             1 / idx
-        with pytest.raises(TypeError, match="cannot perform __floordiv__"):
+
+        floordiv_err = lmsg.replace("*", "//")
+        with pytest.raises(TypeError, match=floordiv_err):
             idx // 1
-        with pytest.raises(TypeError, match="cannot perform __rfloordiv__"):
+        floordiv_err = rmsg.replace("*", "//")
+        with pytest.raises(TypeError, match=floordiv_err):
             1 // idx
 
     def test_logical_compat(self):