Skip to content

TYP/REF: use OpsMixin for arithmetic methods #36994

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
merged 3 commits into from
Oct 10, 2020
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
70 changes: 70 additions & 0 deletions pandas/core/arraylike.py
Original file line number Diff line number Diff line change
Expand Up @@ -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)
27 changes: 7 additions & 20 deletions pandas/core/indexes/base.py
Original file line number Diff line number Diff line change
Expand Up @@ -5408,29 +5408,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.
"""
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):
Expand All @@ -5455,7 +5443,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):
"""
Expand Down
52 changes: 16 additions & 36 deletions pandas/core/indexes/range.py
Original file line number Diff line number Diff line change
Expand Up @@ -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):
Expand All @@ -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()

Expand Down Expand Up @@ -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)
26 changes: 1 addition & 25 deletions pandas/core/ops/__init__.py
Original file line number Diff line number Diff line change
Expand Up @@ -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,
Expand Down Expand Up @@ -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)
Expand Down
93 changes: 53 additions & 40 deletions pandas/core/ops/methods.py
Original file line number Diff line number Diff line change
Expand Up @@ -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,
Expand All @@ -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):
Expand Down Expand Up @@ -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__),
Expand All @@ -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__"]),
Expand Down Expand Up @@ -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
Expand All @@ -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),
Expand Down
Loading