CLN/REF: Unify Arithmetic Methods (#27413)

Author: jbrockmendel

pandas/core/arrays/datetimelike.py

@@ -39,7 +39,7 @@
 from pandas.core.dtypes.missing import is_valid_nat_for_dtype, isna
 
 from pandas._typing import DatetimeLikeScalar
-from pandas.core import missing, nanops
+from pandas.core import missing, nanops, ops
 from pandas.core.algorithms import checked_add_with_arr, take, unique1d, value_counts
 import pandas.core.common as com
 
@@ -926,6 +926,21 @@ def _is_unique(self):
     # ------------------------------------------------------------------
     # Arithmetic Methods
 
+    # pow is invalid for all three subclasses; TimedeltaArray will override
+    #  the multiplication and division ops
+    __pow__ = ops.make_invalid_op("__pow__")
+    __rpow__ = ops.make_invalid_op("__rpow__")
+    __mul__ = ops.make_invalid_op("__mul__")
+    __rmul__ = ops.make_invalid_op("__rmul__")
+    __truediv__ = ops.make_invalid_op("__truediv__")
+    __rtruediv__ = ops.make_invalid_op("__rtruediv__")
+    __floordiv__ = ops.make_invalid_op("__floordiv__")
+    __rfloordiv__ = ops.make_invalid_op("__rfloordiv__")
+    __mod__ = ops.make_invalid_op("__mod__")
+    __rmod__ = ops.make_invalid_op("__rmod__")
+    __divmod__ = ops.make_invalid_op("__divmod__")
+    __rdivmod__ = ops.make_invalid_op("__rdivmod__")
+
     def _add_datetimelike_scalar(self, other):
         # Overriden by TimedeltaArray
         raise TypeError(

pandas/core/indexes/base.py

@@ -9,7 +9,7 @@
 from pandas._libs import algos as libalgos, index as libindex, lib
 import pandas._libs.join as libjoin
 from pandas._libs.lib import is_datetime_array
-from pandas._libs.tslibs import OutOfBoundsDatetime, Timedelta, Timestamp
+from pandas._libs.tslibs import OutOfBoundsDatetime, Timestamp
 from pandas._libs.tslibs.timezones import tz_compare
 from pandas.compat import set_function_name
 from pandas.compat.numpy import function as nv
@@ -55,7 +55,6 @@
     ABCPandasArray,
     ABCPeriodIndex,
     ABCSeries,
-    ABCTimedeltaArray,
     ABCTimedeltaIndex,
 )
 from pandas.core.dtypes.missing import array_equivalent, isna
@@ -126,28 +125,8 @@ def cmp_method(self, other):
 
 def _make_arithmetic_op(op, cls):
     def index_arithmetic_method(self, other):
-        if isinstance(other, (ABCSeries, ABCDataFrame)):
-            return NotImplemented
-        elif isinstance(other, ABCTimedeltaIndex):
-            # Defer to subclass implementation
+        if isinstance(other, (ABCSeries, ABCDataFrame, ABCTimedeltaIndex)):
             return NotImplemented
-        elif isinstance(
-            other, (np.ndarray, ABCTimedeltaArray)
-        ) and is_timedelta64_dtype(other):
-            # GH#22390; wrap in Series for op, this will in turn wrap in
-            # TimedeltaIndex, but will correctly raise TypeError instead of
-            # NullFrequencyError for add/sub ops
-            from pandas import Series
-
-            other = Series(other)
-            out = op(self, other)
-            return Index(out, name=self.name)
-
-        # handle time-based others
-        if isinstance(other, (ABCDateOffset, np.timedelta64, timedelta)):
-            return self._evaluate_with_timedelta_like(other, op)
-
-        other = self._validate_for_numeric_binop(other, op)
 
         from pandas import Series
 
@@ -5332,32 +5311,6 @@ def drop(self, labels, errors="raise"):
     # --------------------------------------------------------------------
     # Generated Arithmetic, Comparison, and Unary Methods
 
-    def _evaluate_with_timedelta_like(self, other, op):
-        # Timedelta knows how to operate with np.array, so dispatch to that
-        # operation and then wrap the results
-        if self._is_numeric_dtype and op.__name__ in ["add", "sub", "radd", "rsub"]:
-            raise TypeError(
-                "Operation {opname} between {cls} and {other} "
-                "is invalid".format(
-                    opname=op.__name__, cls=self.dtype, other=type(other).__name__
-                )
-            )
-
-        other = Timedelta(other)
-        values = self.values
-
-        with np.errstate(all="ignore"):
-            result = op(values, other)
-
-        attrs = self._get_attributes_dict()
-        attrs = self._maybe_update_attributes(attrs)
-        if op == divmod:
-            return Index(result[0], **attrs), Index(result[1], **attrs)
-        return Index(result, **attrs)
-
-    def _evaluate_with_datetime_like(self, other, op):
-        raise TypeError("can only perform ops with datetime like values")
-
     @classmethod
     def _add_comparison_methods(cls):
         """

pandas/core/indexes/datetimelike.py

@@ -62,6 +62,16 @@ def method(self, *args, **kwargs):
     return method
 
 
+def _make_wrapped_arith_op(opname):
+    def method(self, other):
+        meth = getattr(self._data, opname)
+        result = meth(maybe_unwrap_index(other))
+        return wrap_arithmetic_op(self, other, result)
+
+    method.__name__ = opname
+    return method
+
+
 class DatetimeIndexOpsMixin(ExtensionOpsMixin):
     """
     common ops mixin to support a unified interface datetimelike Index
@@ -531,6 +541,19 @@ def __rsub__(self, other):
 
         cls.__rsub__ = __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__")
+
     def isin(self, values, level=None):
         """
         Compute boolean array of whether each index value is found in the

pandas/core/indexes/timedeltas.py

@@ -30,27 +30,13 @@
 from pandas.core.indexes.datetimelike import (
     DatetimeIndexOpsMixin,
     DatetimelikeDelegateMixin,
-    maybe_unwrap_index,
-    wrap_arithmetic_op,
 )
 from pandas.core.indexes.numeric import Int64Index
 from pandas.core.ops import get_op_result_name
 
 from pandas.tseries.frequencies import to_offset
 
 
-def _make_wrapped_arith_op(opname):
-
-    meth = getattr(TimedeltaArray, opname)
-
-    def method(self, other):
-        result = meth(self._data, maybe_unwrap_index(other))
-        return wrap_arithmetic_op(self, other, result)
-
-    method.__name__ = opname
-    return method
-
-
 class TimedeltaDelegateMixin(DatetimelikeDelegateMixin):
     # Most attrs are dispatched via datetimelike_{ops,methods}
     # Some are "raw" methods, the result is not not re-boxed in an Index
@@ -320,17 +306,6 @@ def _format_native_types(self, na_rep="NaT", date_format=None, **kwargs):
     # -------------------------------------------------------------------
     # Wrapping TimedeltaArray
 
-    __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__")
-
     # Compat for frequency inference, see GH#23789
     _is_monotonic_increasing = Index.is_monotonic_increasing
     _is_monotonic_decreasing = Index.is_monotonic_decreasing

pandas/core/ops/__init__.py

@@ -425,8 +425,8 @@ def masked_arith_op(x, y, op):
     # For Series `x` is 1D so ravel() is a no-op; calling it anyway makes
     # the logic valid for both Series and DataFrame ops.
     xrav = x.ravel()
-    assert isinstance(x, (np.ndarray, ABCSeries)), type(x)
-    if isinstance(y, (np.ndarray, ABCSeries, ABCIndexClass)):
+    assert isinstance(x, np.ndarray), type(x)
+    if isinstance(y, np.ndarray):
         dtype = find_common_type([x.dtype, y.dtype])
         result = np.empty(x.size, dtype=dtype)
 
@@ -444,7 +444,7 @@ def masked_arith_op(x, y, op):
 
         if mask.any():
             with np.errstate(all="ignore"):
-                result[mask] = op(xrav[mask], com.values_from_object(yrav[mask]))
+                result[mask] = op(xrav[mask], yrav[mask])
 
     else:
         assert is_scalar(y), type(y)