REF/TYP: use OpsMixin for arithmetic methods (#36994)

Author: jbrockmendel

pandas/core/arraylike.py

@@ -72,3 +72,73 @@ def __xor__(self, other):
     @unpack_zerodim_and_defer("__rxor__")
     def __rxor__(self, other):
         return self._logical_method(other, roperator.rxor)
+
+    # -------------------------------------------------------------
+    # Arithmetic Methods
+
+    def _arith_method(self, other, op):
+        return NotImplemented
+
+    @unpack_zerodim_and_defer("__add__")
+    def __add__(self, other):
+        return self._arith_method(other, operator.add)
+
+    @unpack_zerodim_and_defer("__radd__")
+    def __radd__(self, other):
+        return self._arith_method(other, roperator.radd)
+
+    @unpack_zerodim_and_defer("__sub__")
+    def __sub__(self, other):
+        return self._arith_method(other, operator.sub)
+
+    @unpack_zerodim_and_defer("__rsub__")
+    def __rsub__(self, other):
+        return self._arith_method(other, roperator.rsub)
+
+    @unpack_zerodim_and_defer("__mul__")
+    def __mul__(self, other):
+        return self._arith_method(other, operator.mul)
+
+    @unpack_zerodim_and_defer("__rmul__")
+    def __rmul__(self, other):
+        return self._arith_method(other, roperator.rmul)
+
+    @unpack_zerodim_and_defer("__truediv__")
+    def __truediv__(self, other):
+        return self._arith_method(other, operator.truediv)
+
+    @unpack_zerodim_and_defer("__rtruediv__")
+    def __rtruediv__(self, other):
+        return self._arith_method(other, roperator.rtruediv)
+
+    @unpack_zerodim_and_defer("__floordiv__")
+    def __floordiv__(self, other):
+        return self._arith_method(other, operator.floordiv)
+
+    @unpack_zerodim_and_defer("__rfloordiv")
+    def __rfloordiv__(self, other):
+        return self._arith_method(other, roperator.rfloordiv)
+
+    @unpack_zerodim_and_defer("__mod__")
+    def __mod__(self, other):
+        return self._arith_method(other, operator.mod)
+
+    @unpack_zerodim_and_defer("__rmod__")
+    def __rmod__(self, other):
+        return self._arith_method(other, roperator.rmod)
+
+    @unpack_zerodim_and_defer("__divmod__")
+    def __divmod__(self, other):
+        return self._arith_method(other, divmod)
+
+    @unpack_zerodim_and_defer("__rdivmod__")
+    def __rdivmod__(self, other):
+        return self._arith_method(other, roperator.rdivmod)
+
+    @unpack_zerodim_and_defer("__pow__")
+    def __pow__(self, other):
+        return self._arith_method(other, operator.pow)
+
+    @unpack_zerodim_and_defer("__rpow__")
+    def __rpow__(self, other):
+        return self._arith_method(other, roperator.rpow)

pandas/core/indexes/base.py

@@ -5406,29 +5406,17 @@ def _cmp_method(self, other, op):
 
         return result
 
-    @classmethod
-    def _add_numeric_methods_binary(cls):
+    def _arith_method(self, other, op):
         """
-        Add in numeric methods.
+        Wrapper used to dispatch arithmetic operations.
         """
-        setattr(cls, "__add__", _make_arithmetic_op(operator.add, cls))
-        setattr(cls, "__radd__", _make_arithmetic_op(ops.radd, cls))
-        setattr(cls, "__sub__", _make_arithmetic_op(operator.sub, cls))
-        setattr(cls, "__rsub__", _make_arithmetic_op(ops.rsub, cls))
-        setattr(cls, "__rpow__", _make_arithmetic_op(ops.rpow, cls))
-        setattr(cls, "__pow__", _make_arithmetic_op(operator.pow, cls))
 
-        setattr(cls, "__truediv__", _make_arithmetic_op(operator.truediv, cls))
-        setattr(cls, "__rtruediv__", _make_arithmetic_op(ops.rtruediv, cls))
+        from pandas import Series
 
-        setattr(cls, "__mod__", _make_arithmetic_op(operator.mod, cls))
-        setattr(cls, "__rmod__", _make_arithmetic_op(ops.rmod, cls))
-        setattr(cls, "__floordiv__", _make_arithmetic_op(operator.floordiv, cls))
-        setattr(cls, "__rfloordiv__", _make_arithmetic_op(ops.rfloordiv, cls))
-        setattr(cls, "__divmod__", _make_arithmetic_op(divmod, cls))
-        setattr(cls, "__rdivmod__", _make_arithmetic_op(ops.rdivmod, cls))
-        setattr(cls, "__mul__", _make_arithmetic_op(operator.mul, cls))
-        setattr(cls, "__rmul__", _make_arithmetic_op(ops.rmul, cls))
+        result = op(Series(self), other)
+        if isinstance(result, tuple):
+            return (Index(result[0]), Index(result[1]))
+        return Index(result)
 
     @classmethod
     def _add_numeric_methods_unary(cls):
@@ -5453,7 +5441,6 @@ def _evaluate_numeric_unary(self):
     @classmethod
     def _add_numeric_methods(cls):
         cls._add_numeric_methods_unary()
-        cls._add_numeric_methods_binary()
 
     def any(self, *args, **kwargs):
         """

pandas/core/indexes/range.py

@@ -811,16 +811,13 @@ def any(self, *args, **kwargs) -> bool:
 
     # --------------------------------------------------------------------
 
-    def _arith_method(self, other, op, step=False):
+    def _arith_method(self, other, op):
         """
         Parameters
         ----------
         other : Any
         op : callable that accepts 2 params
             perform the binary op
-        step : callable, optional, default to False
-            op to apply to the step parm if not None
-            if False, use the existing step
         """
 
         if isinstance(other, ABCTimedeltaIndex):
@@ -834,6 +831,21 @@ def _arith_method(self, other, op, step=False):
             # Must be an np.ndarray; GH#22390
             return op(self._int64index, other)
 
+        if op in [
+            operator.pow,
+            ops.rpow,
+            operator.mod,
+            ops.rmod,
+            ops.rfloordiv,
+            divmod,
+            ops.rdivmod,
+        ]:
+            return op(self._int64index, other)
+
+        step = False
+        if op in [operator.mul, ops.rmul, operator.truediv, ops.rtruediv]:
+            step = op
+
         other = extract_array(other, extract_numpy=True)
         attrs = self._get_attributes_dict()
 
@@ -871,35 +883,3 @@ def _arith_method(self, other, op, step=False):
             # Defer to Int64Index implementation
             return op(self._int64index, other)
             # TODO: Do attrs get handled reliably?
-
-    @unpack_zerodim_and_defer("__add__")
-    def __add__(self, other):
-        return self._arith_method(other, operator.add)
-
-    @unpack_zerodim_and_defer("__radd__")
-    def __radd__(self, other):
-        return self._arith_method(other, ops.radd)
-
-    @unpack_zerodim_and_defer("__sub__")
-    def __sub__(self, other):
-        return self._arith_method(other, operator.sub)
-
-    @unpack_zerodim_and_defer("__rsub__")
-    def __rsub__(self, other):
-        return self._arith_method(other, ops.rsub)
-
-    @unpack_zerodim_and_defer("__mul__")
-    def __mul__(self, other):
-        return self._arith_method(other, operator.mul, step=operator.mul)
-
-    @unpack_zerodim_and_defer("__rmul__")
-    def __rmul__(self, other):
-        return self._arith_method(other, ops.rmul, step=ops.rmul)
-
-    @unpack_zerodim_and_defer("__truediv__")
-    def __truediv__(self, other):
-        return self._arith_method(other, operator.truediv, step=operator.truediv)
-
-    @unpack_zerodim_and_defer("__rtruediv__")
-    def __rtruediv__(self, other):
-        return self._arith_method(other, ops.rtruediv, step=ops.rtruediv)

pandas/core/ops/__init__.py

@@ -19,15 +19,14 @@
 from pandas.core.dtypes.missing import isna
 
 from pandas.core import algorithms
-from pandas.core.construction import extract_array
 from pandas.core.ops.array_ops import (  # noqa:F401
     arithmetic_op,
     comp_method_OBJECT_ARRAY,
     comparison_op,
     get_array_op,
     logical_op,
 )
-from pandas.core.ops.common import unpack_zerodim_and_defer
+from pandas.core.ops.common import unpack_zerodim_and_defer  # noqa:F401
 from pandas.core.ops.docstrings import (
     _arith_doc_FRAME,
     _flex_comp_doc_FRAME,
@@ -300,29 +299,6 @@ def align_method_SERIES(left: "Series", right, align_asobject: bool = False):
     return left, right
 
 
-def arith_method_SERIES(cls, op, special):
-    """
-    Wrapper function for Series arithmetic operations, to avoid
-    code duplication.
-    """
-    assert special  # non-special uses flex_method_SERIES
-    op_name = _get_op_name(op, special)
-
-    @unpack_zerodim_and_defer(op_name)
-    def wrapper(left, right):
-        res_name = get_op_result_name(left, right)
-        left, right = align_method_SERIES(left, right)
-
-        lvalues = extract_array(left, extract_numpy=True)
-        rvalues = extract_array(right, extract_numpy=True)
-        result = arithmetic_op(lvalues, rvalues, op)
-
-        return left._construct_result(result, name=res_name)
-
-    wrapper.__name__ = op_name
-    return wrapper
-
-
 def flex_method_SERIES(cls, op, special):
     assert not special  # "special" also means "not flex"
     name = _get_op_name(op, special)

pandas/core/ops/methods.py

@@ -45,7 +45,6 @@ def _get_method_wrappers(cls):
     #  are no longer in __init__
     from pandas.core.ops import (
         arith_method_FRAME,
-        arith_method_SERIES,
         comp_method_FRAME,
         flex_comp_method_FRAME,
         flex_method_SERIES,
@@ -55,7 +54,7 @@ def _get_method_wrappers(cls):
         # Just Series
         arith_flex = flex_method_SERIES
         comp_flex = flex_method_SERIES
-        arith_special = arith_method_SERIES
+        arith_special = None
         comp_special = None
         bool_special = None
     elif issubclass(cls, ABCDataFrame):
@@ -105,20 +104,19 @@ def f(self, other):
         f.__name__ = f"__i{name}__"
         return f
 
-    new_methods.update(
-        dict(
-            __iadd__=_wrap_inplace_method(new_methods["__add__"]),
-            __isub__=_wrap_inplace_method(new_methods["__sub__"]),
-            __imul__=_wrap_inplace_method(new_methods["__mul__"]),
-            __itruediv__=_wrap_inplace_method(new_methods["__truediv__"]),
-            __ifloordiv__=_wrap_inplace_method(new_methods["__floordiv__"]),
-            __imod__=_wrap_inplace_method(new_methods["__mod__"]),
-            __ipow__=_wrap_inplace_method(new_methods["__pow__"]),
-        )
-    )
-
     if bool_method is None:
-        # Series gets bool_method via OpsMixin
+        # Series gets bool_method, arith_method via OpsMixin
+        new_methods.update(
+            dict(
+                __iadd__=_wrap_inplace_method(cls.__add__),
+                __isub__=_wrap_inplace_method(cls.__sub__),
+                __imul__=_wrap_inplace_method(cls.__mul__),
+                __itruediv__=_wrap_inplace_method(cls.__truediv__),
+                __ifloordiv__=_wrap_inplace_method(cls.__floordiv__),
+                __imod__=_wrap_inplace_method(cls.__mod__),
+                __ipow__=_wrap_inplace_method(cls.__pow__),
+            )
+        )
         new_methods.update(
             dict(
                 __iand__=_wrap_inplace_method(cls.__and__),
@@ -127,6 +125,17 @@ def f(self, other):
             )
         )
     else:
+        new_methods.update(
+            dict(
+                __iadd__=_wrap_inplace_method(new_methods["__add__"]),
+                __isub__=_wrap_inplace_method(new_methods["__sub__"]),
+                __imul__=_wrap_inplace_method(new_methods["__mul__"]),
+                __itruediv__=_wrap_inplace_method(new_methods["__truediv__"]),
+                __ifloordiv__=_wrap_inplace_method(new_methods["__floordiv__"]),
+                __imod__=_wrap_inplace_method(new_methods["__mod__"]),
+                __ipow__=_wrap_inplace_method(new_methods["__pow__"]),
+            )
+        )
         new_methods.update(
             dict(
                 __iand__=_wrap_inplace_method(new_methods["__and__"]),
@@ -172,30 +181,34 @@ def _create_methods(cls, arith_method, comp_method, bool_method, special):
     have_divmod = issubclass(cls, ABCSeries)
     # divmod is available for Series
 
-    new_methods = dict(
-        add=arith_method(cls, operator.add, special),
-        radd=arith_method(cls, radd, special),
-        sub=arith_method(cls, operator.sub, special),
-        mul=arith_method(cls, operator.mul, special),
-        truediv=arith_method(cls, operator.truediv, special),
-        floordiv=arith_method(cls, operator.floordiv, special),
-        mod=arith_method(cls, operator.mod, special),
-        pow=arith_method(cls, operator.pow, special),
-        # not entirely sure why this is necessary, but previously was included
-        # so it's here to maintain compatibility
-        rmul=arith_method(cls, rmul, special),
-        rsub=arith_method(cls, rsub, special),
-        rtruediv=arith_method(cls, rtruediv, special),
-        rfloordiv=arith_method(cls, rfloordiv, special),
-        rpow=arith_method(cls, rpow, special),
-        rmod=arith_method(cls, rmod, special),
-    )
-    new_methods["div"] = new_methods["truediv"]
-    new_methods["rdiv"] = new_methods["rtruediv"]
-    if have_divmod:
-        # divmod doesn't have an op that is supported by numexpr
-        new_methods["divmod"] = arith_method(cls, divmod, special)
-        new_methods["rdivmod"] = arith_method(cls, rdivmod, special)
+    new_methods = {}
+    if arith_method is not None:
+        new_methods.update(
+            dict(
+                add=arith_method(cls, operator.add, special),
+                radd=arith_method(cls, radd, special),
+                sub=arith_method(cls, operator.sub, special),
+                mul=arith_method(cls, operator.mul, special),
+                truediv=arith_method(cls, operator.truediv, special),
+                floordiv=arith_method(cls, operator.floordiv, special),
+                mod=arith_method(cls, operator.mod, special),
+                pow=arith_method(cls, operator.pow, special),
+                # not entirely sure why this is necessary, but previously was included
+                # so it's here to maintain compatibility
+                rmul=arith_method(cls, rmul, special),
+                rsub=arith_method(cls, rsub, special),
+                rtruediv=arith_method(cls, rtruediv, special),
+                rfloordiv=arith_method(cls, rfloordiv, special),
+                rpow=arith_method(cls, rpow, special),
+                rmod=arith_method(cls, rmod, special),
+            )
+        )
+        new_methods["div"] = new_methods["truediv"]
+        new_methods["rdiv"] = new_methods["rtruediv"]
+        if have_divmod:
+            # divmod doesn't have an op that is supported by numexpr
+            new_methods["divmod"] = arith_method(cls, divmod, special)
+            new_methods["rdivmod"] = arith_method(cls, rdivmod, special)
 
     if comp_method is not None:
         # Series already has this pinned
@@ -210,7 +223,7 @@ def _create_methods(cls, arith_method, comp_method, bool_method, special):
             )
         )
 
-    if bool_method:
+    if bool_method is not None:
         new_methods.update(
             dict(
                 and_=bool_method(cls, operator.and_, special),