Skip to content

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

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
merged 13 commits into from
Aug 6, 2020
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
2 changes: 1 addition & 1 deletion doc/source/whatsnew/v1.2.0.rst
Original file line number Diff line number Diff line change
Expand Up @@ -17,7 +17,7 @@ Enhancements

Other enhancements
^^^^^^^^^^^^^^^^^^

- :class:`Index` with object dtype supports division and multiplication (:issue:`34160`)
-
-

Expand Down
58 changes: 3 additions & 55 deletions pandas/core/indexes/base.py
Original file line number Diff line number Diff line change
Expand Up @@ -2375,31 +2375,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)

Expand Down Expand Up @@ -5291,38 +5270,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):
"""
Expand All @@ -5338,11 +5285,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)

Expand Down Expand Up @@ -5502,7 +5450,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()

Expand Down
2 changes: 0 additions & 2 deletions pandas/core/indexes/category.py
Original file line number Diff line number Diff line change
Expand Up @@ -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()
1 change: 0 additions & 1 deletion pandas/core/indexes/datetimes.py
Original file line number Diff line number Diff line change
Expand Up @@ -842,7 +842,6 @@ def indexer_between_time(
return mask.nonzero()[0]


DatetimeIndex._add_numeric_methods_disabled()
DatetimeIndex._add_logical_methods_disabled()


Expand Down
35 changes: 35 additions & 0 deletions pandas/core/indexes/multi.py
Original file line number Diff line number Diff line change
Expand Up @@ -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,
Expand Down Expand Up @@ -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()
Expand Down
1 change: 0 additions & 1 deletion pandas/core/indexes/period.py
Original file line number Diff line number Diff line change
Expand Up @@ -724,7 +724,6 @@ def memory_usage(self, deep=False):
return result


PeriodIndex._add_numeric_methods_disabled()
PeriodIndex._add_logical_methods_disabled()


Expand Down
14 changes: 0 additions & 14 deletions pandas/tests/arithmetic/test_numeric.py
Original file line number Diff line number Diff line change
Expand Up @@ -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)])
Expand Down
9 changes: 8 additions & 1 deletion pandas/tests/indexes/categorical/test_category.py
Original file line number Diff line number Diff line change
Expand Up @@ -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)

Expand Down
33 changes: 26 additions & 7 deletions pandas/tests/indexes/common.py
Original file line number Diff line number Diff line change
Expand Up @@ -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):
Expand Down