diff --git a/pandas/core/base.py b/pandas/core/base.py index 048831a74d1a5..b0993bbc619dc 100644 --- a/pandas/core/base.py +++ b/pandas/core/base.py @@ -52,7 +52,10 @@ remove_na_arraylike, ) -from pandas.core import algorithms +from pandas.core import ( + algorithms, + ops, +) from pandas.core.accessor import DirNamesMixin from pandas.core.algorithms import ( duplicated, @@ -61,7 +64,11 @@ ) from pandas.core.arraylike import OpsMixin from pandas.core.arrays import ExtensionArray -from pandas.core.construction import create_series_with_explicit_dtype +from pandas.core.construction import ( + create_series_with_explicit_dtype, + ensure_wrapped_if_datetimelike, + extract_array, +) import pandas.core.nanops as nanops if TYPE_CHECKING: @@ -1238,3 +1245,23 @@ def _duplicated( self, keep: Literal["first", "last", False] = "first" ) -> npt.NDArray[np.bool_]: return duplicated(self._values, keep=keep) + + def _arith_method(self, other, op): + res_name = ops.get_op_result_name(self, other) + + lvalues = self._values + rvalues = extract_array(other, extract_numpy=True, extract_range=True) + rvalues = ops.maybe_prepare_scalar_for_op(rvalues, lvalues.shape) + rvalues = ensure_wrapped_if_datetimelike(rvalues) + + with np.errstate(all="ignore"): + result = ops.arithmetic_op(lvalues, rvalues, op) + + return self._construct_result(result, name=res_name) + + def _construct_result(self, result, name): + """ + Construct an appropriately-wrapped result from the ArrayLike result + of an arithmetic-like operation. + """ + raise AbstractMethodError(self) diff --git a/pandas/core/indexes/base.py b/pandas/core/indexes/base.py index 6c33510ddf1bc..a1c411031a465 100644 --- a/pandas/core/indexes/base.py +++ b/pandas/core/indexes/base.py @@ -6364,7 +6364,10 @@ def _cmp_method(self, other, op): arr[self.isna()] = False return arr elif op in {operator.ne, operator.lt, operator.gt}: - return np.zeros(len(self), dtype=bool) + arr = np.zeros(len(self), dtype=bool) + if self._can_hold_na and not isinstance(self, ABCMultiIndex): + arr[self.isna()] = True + return arr if isinstance(other, (np.ndarray, Index, ABCSeries, ExtensionArray)) and len( self @@ -6381,6 +6384,9 @@ def _cmp_method(self, other, op): with np.errstate(all="ignore"): result = op(self._values, other) + elif isinstance(self._values, ExtensionArray): + result = op(self._values, other) + elif is_object_dtype(self.dtype) and not isinstance(self, ABCMultiIndex): # don't pass MultiIndex with np.errstate(all="ignore"): @@ -6392,17 +6398,26 @@ def _cmp_method(self, other, op): return result - def _arith_method(self, other, op): - """ - Wrapper used to dispatch arithmetic operations. - """ + def _construct_result(self, result, name): + if isinstance(result, tuple): + return ( + Index._with_infer(result[0], name=name), + Index._with_infer(result[1], name=name), + ) + return Index._with_infer(result, name=name) - from pandas import Series + def _arith_method(self, other, op): + 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 - result = op(Series(self), other) - if isinstance(result, tuple): - return (Index._with_infer(result[0]), Index(result[1])) - return Index._with_infer(result) + return super()._arith_method(other, op) @final def _unary_method(self, op): diff --git a/pandas/core/indexes/extension.py b/pandas/core/indexes/extension.py index cea67aa13edd1..44267895a989e 100644 --- a/pandas/core/indexes/extension.py +++ b/pandas/core/indexes/extension.py @@ -26,13 +26,9 @@ from pandas.core.dtypes.common import ( is_dtype_equal, - is_object_dtype, pandas_dtype, ) -from pandas.core.dtypes.generic import ( - ABCDataFrame, - ABCSeries, -) +from pandas.core.dtypes.generic import ABCDataFrame from pandas.core.arrays import ( Categorical, @@ -45,7 +41,6 @@ from pandas.core.arrays.base import ExtensionArray from pandas.core.indexers import deprecate_ndim_indexing from pandas.core.indexes.base import Index -from pandas.core.ops import get_op_result_name if TYPE_CHECKING: @@ -154,94 +149,6 @@ def wrapper(cls): return wrapper -def _make_wrapped_comparison_op(opname: str): - """ - Create a comparison method that dispatches to ``._data``. - """ - - def wrapper(self, other): - if isinstance(other, ABCSeries): - # the arrays defer to Series for comparison ops but the indexes - # don't, so we have to unwrap here. - other = other._values - - other = _maybe_unwrap_index(other) - - op = getattr(self._data, opname) - return op(other) - - wrapper.__name__ = opname - return wrapper - - -def _make_wrapped_arith_op(opname: str): - 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 - - try: - meth = getattr(self._data, opname) - except AttributeError as err: - # e.g. Categorical, IntervalArray - cls = type(self).__name__ - raise TypeError( - f"cannot perform {opname} with this index type: {cls}" - ) from err - - result = meth(_maybe_unwrap_index(other)) - return _wrap_arithmetic_op(self, other, result) - - method.__name__ = opname - return method - - -def _wrap_arithmetic_op(self, other, result): - if result is NotImplemented: - return NotImplemented - - if isinstance(result, tuple): - # divmod, rdivmod - assert len(result) == 2 - return ( - _wrap_arithmetic_op(self, other, result[0]), - _wrap_arithmetic_op(self, other, result[1]), - ) - - if not isinstance(result, Index): - # Index.__new__ will choose appropriate subclass for dtype - result = Index(result) - - res_name = get_op_result_name(self, other) - result.name = res_name - return result - - -def _maybe_unwrap_index(obj): - """ - If operating against another Index object, we need to unwrap the underlying - data before deferring to the DatetimeArray/TimedeltaArray/PeriodArray - implementation, otherwise we will incorrectly return NotImplemented. - - Parameters - ---------- - obj : object - - Returns - ------- - unwrapped object - """ - if isinstance(obj, Index): - return obj._data - return obj - - class ExtensionIndex(Index): """ Index subclass for indexes backed by ExtensionArray. @@ -284,30 +191,6 @@ def _simple_new( result._reset_identity() return result - __eq__ = _make_wrapped_comparison_op("__eq__") - __ne__ = _make_wrapped_comparison_op("__ne__") - __lt__ = _make_wrapped_comparison_op("__lt__") - __gt__ = _make_wrapped_comparison_op("__gt__") - __le__ = _make_wrapped_comparison_op("__le__") - __ge__ = _make_wrapped_comparison_op("__ge__") - - __add__ = _make_wrapped_arith_op("__add__") - __sub__ = _make_wrapped_arith_op("__sub__") - __radd__ = _make_wrapped_arith_op("__radd__") - __rsub__ = _make_wrapped_arith_op("__rsub__") - __pow__ = _make_wrapped_arith_op("__pow__") - __rpow__ = _make_wrapped_arith_op("__rpow__") - __mul__ = _make_wrapped_arith_op("__mul__") - __rmul__ = _make_wrapped_arith_op("__rmul__") - __floordiv__ = _make_wrapped_arith_op("__floordiv__") - __rfloordiv__ = _make_wrapped_arith_op("__rfloordiv__") - __mod__ = _make_wrapped_arith_op("__mod__") - __rmod__ = _make_wrapped_arith_op("__rmod__") - __divmod__ = _make_wrapped_arith_op("__divmod__") - __rdivmod__ = _make_wrapped_arith_op("__rdivmod__") - __truediv__ = _make_wrapped_arith_op("__truediv__") - __rtruediv__ = _make_wrapped_arith_op("__rtruediv__") - # --------------------------------------------------------------------- # NDarray-Like Methods diff --git a/pandas/core/series.py b/pandas/core/series.py index 6f964ab09e978..c7756e7f42337 100644 --- a/pandas/core/series.py +++ b/pandas/core/series.py @@ -104,7 +104,6 @@ import pandas.core.common as com from pandas.core.construction import ( create_series_with_explicit_dtype, - ensure_wrapped_if_datetimelike, extract_array, is_empty_data, sanitize_array, @@ -5515,18 +5514,8 @@ def _logical_method(self, other, op): return self._construct_result(res_values, name=res_name) def _arith_method(self, other, op): - res_name = ops.get_op_result_name(self, other) self, other = ops.align_method_SERIES(self, other) - - lvalues = self._values - rvalues = extract_array(other, extract_numpy=True, extract_range=True) - rvalues = ops.maybe_prepare_scalar_for_op(rvalues, lvalues.shape) - rvalues = ensure_wrapped_if_datetimelike(rvalues) - - with np.errstate(all="ignore"): - result = ops.arithmetic_op(lvalues, rvalues, op) - - return self._construct_result(result, name=res_name) + return base.IndexOpsMixin._arith_method(self, other, op) Series._add_numeric_operations() diff --git a/pandas/tests/arithmetic/test_datetime64.py b/pandas/tests/arithmetic/test_datetime64.py index 6b3309ba8ea1b..c0f38a1181026 100644 --- a/pandas/tests/arithmetic/test_datetime64.py +++ b/pandas/tests/arithmetic/test_datetime64.py @@ -2146,7 +2146,7 @@ def test_dti_sub_tdi(self, tz_naive_fixture): result = dti - tdi.values tm.assert_index_equal(result, expected) - msg = "cannot subtract DatetimeArray from" + msg = "cannot subtract a datelike from a TimedeltaArray" with pytest.raises(TypeError, match=msg): tdi.values - dti @@ -2172,13 +2172,7 @@ def test_dti_isub_tdi(self, tz_naive_fixture): result -= tdi.values tm.assert_index_equal(result, expected) - msg = "|".join( - [ - "cannot perform __neg__ with this index type:", - "ufunc subtract cannot use operands with types", - "cannot subtract DatetimeArray from", - ] - ) + msg = "cannot subtract a datelike from a TimedeltaArray" with pytest.raises(TypeError, match=msg): tdi.values -= dti diff --git a/pandas/tests/arithmetic/test_period.py b/pandas/tests/arithmetic/test_period.py index 5f93442cae4f6..7d215c940c031 100644 --- a/pandas/tests/arithmetic/test_period.py +++ b/pandas/tests/arithmetic/test_period.py @@ -753,7 +753,7 @@ def test_pi_add_sub_td64_array_non_tick_raises(self): with pytest.raises(TypeError, match=msg): rng - tdarr - msg = r"cannot subtract PeriodArray from timedelta64\[ns\]" + msg = r"cannot subtract period\[Q-DEC\]-dtype from TimedeltaArray" with pytest.raises(TypeError, match=msg): tdarr - rng