REF: call ensure_wrapped_if_datetimelike before the array arithmetic_op

Author: jorisvandenbossche

pandas/core/arrays/numpy_.py

@@ -23,6 +23,7 @@
 )
 from pandas.core.arraylike import OpsMixin
 from pandas.core.arrays._mixins import NDArrayBackedExtensionArray
+from pandas.core.construction import ensure_wrapped_if_datetimelike
 from pandas.core.strings.object_array import ObjectStringArrayMixin
 
 
@@ -395,6 +396,7 @@ def _cmp_method(self, other, op):
             other = other._ndarray
 
         pd_op = ops.get_array_op(op)
+        other = ensure_wrapped_if_datetimelike(other)
         with np.errstate(all="ignore"):
             result = pd_op(self._ndarray, other)
 

pandas/core/internals/ops.py

@@ -8,6 +8,8 @@
 
 from pandas._typing import ArrayLike
 
+from pandas.core.construction import ensure_wrapped_if_datetimelike
+
 if TYPE_CHECKING:
     from pandas.core.internals.blocks import Block
     from pandas.core.internals.managers import BlockManager
@@ -54,6 +56,8 @@ def operate_blockwise(
 
     res_blks: list[Block] = []
     for lvals, rvals, locs, left_ea, right_ea, rblk in _iter_block_pairs(left, right):
+        lvals = ensure_wrapped_if_datetimelike(lvals)
+        rvals = ensure_wrapped_if_datetimelike(rvals)
         res_values = array_op(lvals, rvals)
         if left_ea and not right_ea and hasattr(res_values, "reshape"):
             res_values = res_values.reshape(1, -1)

pandas/core/ops/__init__.py

@@ -29,6 +29,7 @@
     algorithms,
     roperator,
 )
+from pandas.core.construction import ensure_wrapped_if_datetimelike
 from pandas.core.ops.array_ops import (  # noqa:F401
     arithmetic_op,
     comp_method_OBJECT_ARRAY,
@@ -140,6 +141,8 @@ def fill_binop(left, right, fill_value):
             right = right.copy()
             right[right_mask & mask] = fill_value
 
+    left = ensure_wrapped_if_datetimelike(left)
+    right = ensure_wrapped_if_datetimelike(right)
     return left, right
 
 

pandas/core/ops/array_ops.py

@@ -195,20 +195,17 @@ def arithmetic_op(left: ArrayLike, right: Any, op):
         Or a 2-tuple of these in the case of divmod or rdivmod.
     """
 
-    # NB: We assume that extract_array has already been called
-    #  on `left` and `right`.
+    # NB: We assume that extract_array and ensure_wrapped_if_datetimelike
+    #  has already been called on `left` and `right`.
     # We need to special-case datetime64/timedelta64 dtypes (e.g. because numpy
     # casts integer dtypes to timedelta64 when operating with timedelta64 - GH#22390)
-    lvalues = ensure_wrapped_if_datetimelike(left)
-    rvalues = ensure_wrapped_if_datetimelike(right)
-    rvalues = _maybe_upcast_for_op(rvalues, lvalues.shape)
+    right = _maybe_upcast_for_op(right, left.shape)
 
-    if should_extension_dispatch(lvalues, rvalues) or isinstance(rvalues, Timedelta):
+    if should_extension_dispatch(left, right) or isinstance(right, Timedelta):
         # Timedelta is included because numexpr will fail on it, see GH#31457
-        res_values = op(lvalues, rvalues)
-
+        res_values = op(left, right)
     else:
-        res_values = _na_arithmetic_op(lvalues, rvalues, op)
+        res_values = _na_arithmetic_op(left, right, op)
 
     return res_values
 

pandas/core/series.py

@@ -100,6 +100,7 @@
 import pandas.core.common as com
 from pandas.core.construction import (
     create_series_with_explicit_dtype,
+    ensure_wrapped_if_datetimelike,
     extract_array,
     is_empty_data,
     sanitize_array,
@@ -5313,6 +5314,8 @@ def _arith_method(self, other, op):
 
         lvalues = self._values
         rvalues = extract_array(other, extract_numpy=True, extract_range=True)
+        lvalues = ensure_wrapped_if_datetimelike(lvalues)
+        rvalues = ensure_wrapped_if_datetimelike(rvalues)
 
         with np.errstate(all="ignore"):
             result = ops.arithmetic_op(lvalues, rvalues, op)