Skip to content

CLN/REF: Unify Arithmetic Methods #27413

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 6 commits into from
Jul 20, 2019
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
17 changes: 16 additions & 1 deletion pandas/core/arrays/datetimelike.py
Original file line number Diff line number Diff line change
Expand Up @@ -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

Expand Down Expand Up @@ -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(
Expand Down
51 changes: 2 additions & 49 deletions pandas/core/indexes/base.py
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -55,7 +55,6 @@
ABCPandasArray,
ABCPeriodIndex,
ABCSeries,
ABCTimedeltaArray,
ABCTimedeltaIndex,
)
from pandas.core.dtypes.missing import array_equivalent, isna
Expand Down Expand Up @@ -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

Expand Down Expand Up @@ -5336,32 +5315,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):
"""
Expand Down
23 changes: 23 additions & 0 deletions pandas/core/indexes/datetimelike.py
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -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
Expand Down
25 changes: 0 additions & 25 deletions pandas/core/indexes/timedeltas.py
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -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
Expand Down
6 changes: 3 additions & 3 deletions pandas/core/ops/__init__.py
Original file line number Diff line number Diff line change
Expand Up @@ -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)

Expand All @@ -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)
Expand Down