Skip to content

Backport PR #31136 on branch 1.0.x (COMPAT: Return NotImplemented for subclassing) #31405

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
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
15 changes: 14 additions & 1 deletion pandas/core/indexes/extension.py
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -129,6 +133,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)
Expand Down
46 changes: 46 additions & 0 deletions pandas/tests/arithmetic/test_object.py
Original file line number Diff line number Diff line change
@@ -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

Expand Down Expand Up @@ -317,3 +318,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