Skip to content

REF: combine dispatch_to_index_op into dispatch_to_extension_op #27747

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 11 commits into from
Aug 5, 2019
90 changes: 31 additions & 59 deletions pandas/core/ops/__init__.py
Original file line number Diff line number Diff line change
Expand Up @@ -37,6 +37,7 @@
from pandas.core.dtypes.generic import (
ABCDataFrame,
ABCDatetimeArray,
ABCDatetimeIndex,
ABCIndex,
ABCIndexClass,
ABCSeries,
Expand Down Expand Up @@ -90,7 +91,7 @@ def get_op_result_name(left, right):
name : object
Usually a string
"""
# `left` is always a pd.Series when called from within ops
# `left` is always a Series when called from within ops
if isinstance(right, (ABCSeries, ABCIndexClass)):
name = _maybe_match_name(left, right)
else:
Expand Down Expand Up @@ -609,42 +610,6 @@ def column_op(a, b):
return result


def dispatch_to_index_op(op, left, right, index_class):
"""
Wrap Series left in the given index_class to delegate the operation op
to the index implementation. DatetimeIndex and TimedeltaIndex perform
type checking, timezone handling, overflow checks, etc.

Parameters
----------
op : binary operator (operator.add, operator.sub, ...)
left : Series
right : object
index_class : DatetimeIndex or TimedeltaIndex

Returns
-------
result : object, usually DatetimeIndex, TimedeltaIndex, or Series
"""
left_idx = index_class(left)

# avoid accidentally allowing integer add/sub. For datetime64[tz] dtypes,
# left_idx may inherit a freq from a cached DatetimeIndex.
# See discussion in GH#19147.
if getattr(left_idx, "freq", None) is not None:
left_idx = left_idx._shallow_copy(freq=None)
try:
result = op(left_idx, right)
except NullFrequencyError:
# DatetimeIndex and TimedeltaIndex with freq == None raise ValueError
# on add/sub of integers (or int-like). We re-raise as a TypeError.
raise TypeError(
"incompatible type for a datetime/timedelta "
"operation [{name}]".format(name=op.__name__)
)
return result


def dispatch_to_extension_op(op, left, right):
"""
Assume that left or right is a Series backed by an ExtensionArray,
Expand All @@ -665,13 +630,16 @@ def dispatch_to_extension_op(op, left, right):
else:
new_right = right

res_values = op(new_left, new_right)
res_name = get_op_result_name(left, right)

if op.__name__ in ["divmod", "rdivmod"]:
return _construct_divmod_result(left, res_values, left.index, res_name)

return _construct_result(left, res_values, left.index, res_name)
try:
res_values = op(new_left, new_right)
except NullFrequencyError:
# DatetimeIndex and TimedeltaIndex with freq == None raise ValueError
# on add/sub of integers (or int-like). We re-raise as a TypeError.
raise TypeError(
"incompatible type for a datetime/timedelta "
"operation [{name}]".format(name=op.__name__)
)
return res_values


# -----------------------------------------------------------------------------
Expand Down Expand Up @@ -993,22 +961,22 @@ def wrapper(left, right):
)

elif is_datetime64_dtype(left) or is_datetime64tz_dtype(left):
# Give dispatch_to_index_op a chance for tests like
# test_dt64_series_add_intlike, which the index dispatching handles
# specifically.
result = dispatch_to_index_op(op, left, right, pd.DatetimeIndex)
return construct_result(
left, result, index=left.index, name=res_name, dtype=result.dtype
)
from pandas.core.arrays import DatetimeArray

result = dispatch_to_extension_op(op, DatetimeArray(left), right)
return construct_result(left, result, index=left.index, name=res_name)

elif is_extension_array_dtype(left) or (
is_extension_array_dtype(right) and not is_scalar(right)
):
# GH#22378 disallow scalar to exclude e.g. "category", "Int64"
return dispatch_to_extension_op(op, left, right)
result = dispatch_to_extension_op(op, left, right)
return construct_result(left, result, index=left.index, name=res_name)

elif is_timedelta64_dtype(left):
result = dispatch_to_index_op(op, left, right, pd.TimedeltaIndex)
from pandas.core.arrays import TimedeltaArray

result = dispatch_to_extension_op(op, TimedeltaArray(left), right)
return construct_result(left, result, index=left.index, name=res_name)

elif is_timedelta64_dtype(right):
Expand All @@ -1022,7 +990,7 @@ def wrapper(left, right):
# does inference in the case where `result` has object-dtype.
return construct_result(left, result, index=left.index, name=res_name)

elif isinstance(right, (ABCDatetimeArray, pd.DatetimeIndex)):
elif isinstance(right, (ABCDatetimeArray, ABCDatetimeIndex)):
result = op(left._values, right)
return construct_result(left, result, index=left.index, name=res_name)

Expand Down Expand Up @@ -1129,28 +1097,32 @@ def wrapper(self, other, axis=None):
raise ValueError("Can only compare identically-labeled Series objects")

elif is_categorical_dtype(self):
# Dispatch to Categorical implementation; pd.CategoricalIndex
# Dispatch to Categorical implementation; CategoricalIndex
# behavior is non-canonical GH#19513
res_values = dispatch_to_index_op(op, self, other, pd.Categorical)
res_values = dispatch_to_extension_op(op, self, other)
return self._constructor(res_values, index=self.index, name=res_name)

elif is_datetime64_dtype(self) or is_datetime64tz_dtype(self):
# Dispatch to DatetimeIndex to ensure identical
# Series/Index behavior
from pandas.core.arrays import DatetimeArray

res_values = dispatch_to_index_op(op, self, other, pd.DatetimeIndex)
res_values = dispatch_to_extension_op(op, DatetimeArray(self), other)
return self._constructor(res_values, index=self.index, name=res_name)

elif is_timedelta64_dtype(self):
res_values = dispatch_to_index_op(op, self, other, pd.TimedeltaIndex)
from pandas.core.arrays import TimedeltaArray

res_values = dispatch_to_extension_op(op, TimedeltaArray(self), other)
return self._constructor(res_values, index=self.index, name=res_name)

elif is_extension_array_dtype(self) or (
is_extension_array_dtype(other) and not is_scalar(other)
):
# Note: the `not is_scalar(other)` condition rules out
# e.g. other == "category"
return dispatch_to_extension_op(op, self, other)
res_values = dispatch_to_extension_op(op, self, other)
return self._constructor(res_values, index=self.index).rename(res_name)

elif isinstance(other, ABCSeries):
# By this point we have checked that self._indexed_same(other)
Expand Down