REF: combine dispatch_to_index_op into dispatch_to_extension_op (pandas-dev#27747)

Author: jbrockmendel

pandas/core/ops/__init__.py

@@ -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:
@@ -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,
@@ -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
 
 
 # -----------------------------------------------------------------------------
@@ -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):
@@ -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)