From f1bb6a059bc04cbb7833129c39dbcc706f3ed4bb Mon Sep 17 00:00:00 2001 From: Brock Date: Thu, 8 Oct 2020 15:43:20 -0700 Subject: [PATCH] REF/TYP: use OpsMixin for arithmetic methods --- pandas/core/arraylike.py | 70 ++++++++++++++++ pandas/core/indexes/base.py | 27 ++---- pandas/core/indexes/range.py | 52 ++++-------- pandas/core/ops/__init__.py | 26 +----- pandas/core/ops/methods.py | 93 ++++++++++++--------- pandas/core/series.py | 18 ++++ pandas/tests/arithmetic/test_timedelta64.py | 2 +- 7 files changed, 166 insertions(+), 122 deletions(-) diff --git a/pandas/core/arraylike.py b/pandas/core/arraylike.py index 185e9197e01fe..553649212aa5f 100644 --- a/pandas/core/arraylike.py +++ b/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) diff --git a/pandas/core/indexes/base.py b/pandas/core/indexes/base.py index 539f5515a2f8b..567115edb02eb 100644 --- a/pandas/core/indexes/base.py +++ b/pandas/core/indexes/base.py @@ -5404,29 +5404,17 @@ def _cmp_method(self, other, op): return result return ops.invalid_comparison(self, other, op) - @classmethod - def _add_numeric_methods_binary(cls): + def _arith_method(self, other, op): """ - Add in numeric methods. + Wrapper used to dispatch arithmetic operations. """ - cls.__add__ = _make_arithmetic_op(operator.add, cls) - cls.__radd__ = _make_arithmetic_op(ops.radd, cls) - cls.__sub__ = _make_arithmetic_op(operator.sub, cls) - cls.__rsub__ = _make_arithmetic_op(ops.rsub, cls) - cls.__rpow__ = _make_arithmetic_op(ops.rpow, cls) - cls.__pow__ = _make_arithmetic_op(operator.pow, cls) - cls.__truediv__ = _make_arithmetic_op(operator.truediv, cls) - cls.__rtruediv__ = _make_arithmetic_op(ops.rtruediv, cls) + from pandas import Series - cls.__mod__ = _make_arithmetic_op(operator.mod, cls) - cls.__rmod__ = _make_arithmetic_op(ops.rmod, cls) - cls.__floordiv__ = _make_arithmetic_op(operator.floordiv, cls) - cls.__rfloordiv__ = _make_arithmetic_op(ops.rfloordiv, cls) - cls.__divmod__ = _make_arithmetic_op(divmod, cls) - cls.__rdivmod__ = _make_arithmetic_op(ops.rdivmod, cls) - cls.__mul__ = _make_arithmetic_op(operator.mul, cls) - 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): @@ -5451,7 +5439,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): """ diff --git a/pandas/core/indexes/range.py b/pandas/core/indexes/range.py index 4a6bb11bda400..14098ddadb8e2 100644 --- a/pandas/core/indexes/range.py +++ b/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) diff --git a/pandas/core/ops/__init__.py b/pandas/core/ops/__init__.py index ae21f13ea3f49..0de842e8575af 100644 --- a/pandas/core/ops/__init__.py +++ b/pandas/core/ops/__init__.py @@ -19,7 +19,6 @@ 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, @@ -27,7 +26,7 @@ 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) diff --git a/pandas/core/ops/methods.py b/pandas/core/ops/methods.py index 70fd814423c7f..05da378f8964d 100644 --- a/pandas/core/ops/methods.py +++ b/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), diff --git a/pandas/core/series.py b/pandas/core/series.py index 9bd41ca0e76db..fe73a85499c08 100644 --- a/pandas/core/series.py +++ b/pandas/core/series.py @@ -4990,6 +4990,24 @@ def _logical_method(self, other, op): res_values = ops.logical_op(lvalues, rvalues, 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 = extract_array(self, extract_numpy=True) + rvalues = extract_array(other, extract_numpy=True) + result = ops.arithmetic_op(lvalues, rvalues, op) + + return self._construct_result(result, name=res_name) + + def __div__(self, other): + # Alias for backward compat + return self.__truediv__(other) + + def __rdiv__(self, other): + # Alias for backward compat + return self.__rtruediv__(other) + Series._add_numeric_operations() diff --git a/pandas/tests/arithmetic/test_timedelta64.py b/pandas/tests/arithmetic/test_timedelta64.py index b3dfb5d015ab4..3e979aed0551f 100644 --- a/pandas/tests/arithmetic/test_timedelta64.py +++ b/pandas/tests/arithmetic/test_timedelta64.py @@ -2159,7 +2159,7 @@ def test_float_series_rdiv_td64arr(self, box_with_array, names): tdi = tm.box_expected(tdi, box) expected = tm.box_expected(expected, xbox) - result = ser.__rdiv__(tdi) + result = ser.__rtruediv__(tdi) if box is pd.DataFrame: # TODO: Should we skip this case sooner or test something else? assert result is NotImplemented