Skip to content

Commit

Permalink
REF: combine dispatch_to_index_op into dispatch_to_extension_op (#27747)
Browse files Browse the repository at this point in the history
  • Loading branch information
jbrockmendel authored and jreback committed Aug 5, 2019
1 parent ce357d9 commit 9fe8a0f
Showing 1 changed file with 29 additions and 58 deletions.
87 changes: 29 additions & 58 deletions pandas/core/ops/__init__.py
Original file line number Diff line number Diff line change
Expand Up @@ -91,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 @@ -610,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 @@ -666,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 @@ -994,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 Down Expand Up @@ -1130,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

0 comments on commit 9fe8a0f

Please sign in to comment.