diff --git a/pandas/core/indexes/extension.py b/pandas/core/indexes/extension.py index 9ddc5c01030b1..81afb00c6dacc 100644 --- a/pandas/core/indexes/extension.py +++ b/pandas/core/indexes/extension.py @@ -8,7 +8,11 @@ from pandas.compat.numpy import function as nv from pandas.util._decorators import Appender, cache_readonly -from pandas.core.dtypes.common import ensure_platform_int, is_dtype_equal +from pandas.core.dtypes.common import ( + ensure_platform_int, + is_dtype_equal, + is_object_dtype, +) from pandas.core.dtypes.generic import ABCSeries from pandas.core.arrays import ExtensionArray @@ -110,6 +114,15 @@ def wrapper(self, other): def make_wrapped_arith_op(opname): def method(self, other): + if ( + isinstance(other, Index) + and is_object_dtype(other.dtype) + and type(other) is not Index + ): + # We return NotImplemented for object-dtype index *subclasses* so they have + # a chance to implement ops before we unwrap them. + # See https://github.com/pandas-dev/pandas/issues/31109 + return NotImplemented meth = getattr(self._data, opname) result = meth(_maybe_unwrap_index(other)) return _wrap_arithmetic_op(self, other, result) diff --git a/pandas/tests/arithmetic/test_object.py b/pandas/tests/arithmetic/test_object.py index c0d3c9d4977bd..a729715f11b2a 100644 --- a/pandas/tests/arithmetic/test_object.py +++ b/pandas/tests/arithmetic/test_object.py @@ -1,6 +1,7 @@ # Arithmetic tests for DataFrame/Series/Index/Array classes that should # behave identically. # Specifically for object dtype +import datetime from decimal import Decimal import operator @@ -323,3 +324,48 @@ def test_rsub_object(self): with pytest.raises(TypeError): np.array([True, pd.Timestamp.now()]) - index + + +class MyIndex(pd.Index): + # Simple index subclass that tracks ops calls. + + _calls: int + + @classmethod + def _simple_new(cls, values, name=None, dtype=None): + result = object.__new__(cls) + result._data = values + result._index_data = values + result._name = name + result._calls = 0 + + return result._reset_identity() + + def __add__(self, other): + self._calls += 1 + return self._simple_new(self._index_data) + + def __radd__(self, other): + return self.__add__(other) + + +@pytest.mark.parametrize( + "other", + [ + [datetime.timedelta(1), datetime.timedelta(2)], + [datetime.datetime(2000, 1, 1), datetime.datetime(2000, 1, 2)], + [pd.Period("2000"), pd.Period("2001")], + ["a", "b"], + ], + ids=["timedelta", "datetime", "period", "object"], +) +def test_index_ops_defer_to_unknown_subclasses(other): + # https://github.com/pandas-dev/pandas/issues/31109 + values = np.array( + [datetime.date(2000, 1, 1), datetime.date(2000, 1, 2)], dtype=object + ) + a = MyIndex._simple_new(values) + other = pd.Index(other) + result = other + a + assert isinstance(result, MyIndex) + assert a._calls == 1