From 74c965c6aa63504ff45ecb6299ddcea3049e2c21 Mon Sep 17 00:00:00 2001 From: jbrockmendel Date: Wed, 13 May 2020 10:32:10 -0700 Subject: [PATCH 1/2] ENH: enable mul, div on Index by dispatching to Series --- pandas/core/indexes/base.py | 58 +------------------ pandas/core/indexes/category.py | 2 - pandas/core/indexes/datetimes.py | 1 - pandas/core/indexes/multi.py | 35 +++++++++++ pandas/core/indexes/period.py | 1 - pandas/tests/arithmetic/test_numeric.py | 14 ----- .../indexes/categorical/test_category.py | 9 ++- pandas/tests/indexes/common.py | 33 ++++++++--- 8 files changed, 72 insertions(+), 81 deletions(-) diff --git a/pandas/core/indexes/base.py b/pandas/core/indexes/base.py index b8a9827b5effd..20a473047417b 100644 --- a/pandas/core/indexes/base.py +++ b/pandas/core/indexes/base.py @@ -2324,31 +2324,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) @@ -5261,38 +5240,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): """ @@ -5308,11 +5255,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) @@ -5472,7 +5420,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() diff --git a/pandas/core/indexes/category.py b/pandas/core/indexes/category.py index 25df4a0bee737..05d6d35290f0f 100644 --- a/pandas/core/indexes/category.py +++ b/pandas/core/indexes/category.py @@ -768,6 +768,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() diff --git a/pandas/core/indexes/datetimes.py b/pandas/core/indexes/datetimes.py index 5a89c45a3e425..4d51fc8ee276a 100644 --- a/pandas/core/indexes/datetimes.py +++ b/pandas/core/indexes/datetimes.py @@ -804,7 +804,6 @@ def indexer_between_time( return mask.nonzero()[0] -DatetimeIndex._add_numeric_methods_disabled() DatetimeIndex._add_logical_methods_disabled() diff --git a/pandas/core/indexes/multi.py b/pandas/core/indexes/multi.py index f1e1ebcaca1c4..95f14573ee8fc 100644 --- a/pandas/core/indexes/multi.py +++ b/pandas/core/indexes/multi.py @@ -54,6 +54,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, @@ -3598,6 +3599,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() diff --git a/pandas/core/indexes/period.py b/pandas/core/indexes/period.py index b0b85f69396ba..74f117e7b1c52 100644 --- a/pandas/core/indexes/period.py +++ b/pandas/core/indexes/period.py @@ -720,7 +720,6 @@ def memory_usage(self, deep=False): return result -PeriodIndex._add_numeric_methods_disabled() PeriodIndex._add_logical_methods_disabled() diff --git a/pandas/tests/arithmetic/test_numeric.py b/pandas/tests/arithmetic/test_numeric.py index 269235b943e46..3fa0ab2b8aa5e 100644 --- a/pandas/tests/arithmetic/test_numeric.py +++ b/pandas/tests/arithmetic/test_numeric.py @@ -544,20 +544,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)]) diff --git a/pandas/tests/indexes/categorical/test_category.py b/pandas/tests/indexes/categorical/test_category.py index 9765c77c6b60c..c49af917baaec 100644 --- a/pandas/tests/indexes/categorical/test_category.py +++ b/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) diff --git a/pandas/tests/indexes/common.py b/pandas/tests/indexes/common.py index 0f9509c372bdf..a7e1cf3415a31 100644 --- a/pandas/tests/indexes/common.py +++ b/pandas/tests/indexes/common.py @@ -145,22 +145,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): From f32abf9d8a3fbd3352031c152751867561226994 Mon Sep 17 00:00:00 2001 From: Brock Date: Tue, 4 Aug 2020 12:09:18 -0700 Subject: [PATCH 2/2] whatsnew --- doc/source/whatsnew/v1.2.0.rst | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/doc/source/whatsnew/v1.2.0.rst b/doc/source/whatsnew/v1.2.0.rst index b16ca0a80c5b4..cb5f43aa8d00a 100644 --- a/doc/source/whatsnew/v1.2.0.rst +++ b/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`) - -