REF: call ensure_wrapped_if_datetimelike before the array arithmetic_op

Author: jorisvandenbossche

pandas/core/arrays/numpy_.py

@@ -16,6 +16,7 @@
 from pandas.core import nanops, ops
 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
 
 
@@ -348,6 +349,7 @@ def _cmp_method(self, other, op):
             other = other._ndarray
 
         pd_op = ops.get_array_op(op)
+        other = ensure_wrapped_if_datetimelike(other)
         result = pd_op(self._ndarray, other)
 
         if op is divmod or op is ops.rdivmod:

pandas/core/internals/managers.py

@@ -38,7 +38,7 @@
 from pandas.core import ops
 import pandas.core.algorithms as algos
 from pandas.core.arrays.sparse import SparseDtype
-from pandas.core.construction import extract_array
+from pandas.core.construction import ensure_wrapped_if_datetimelike, extract_array
 from pandas.core.indexers import maybe_convert_indices
 from pandas.core.indexes.api import Float64Index, Index, ensure_index
 from pandas.core.internals.base import DataManager
@@ -382,7 +382,20 @@ def operate_scalar(self, other: Scalar, op) -> BlockManager:
         """
         # Get the appropriate array-op to apply to each column/block's values.
         array_op = ops.get_array_op(op)
-        return self.apply(array_op, right=other)
+        result_blocks: List[Block] = []
+
+        for b in self.blocks:
+            left = ensure_wrapped_if_datetimelike(b.values)
+            with np.errstate(all="ignore"):
+                result = array_op(left, right=other)
+
+            applied = b._split_op_result(result)
+            result_blocks = extend_blocks(applied, result_blocks)
+
+        if len(result_blocks) == 0:
+            return self.make_empty(self.axes)
+
+        return type(self).from_blocks(result_blocks, self.axes)
 
     def operate_array(self, other: ArrayLike, op, axis: int) -> BlockManager:
         """

pandas/core/internals/ops.py

@@ -7,6 +7,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
@@ -53,6 +55,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

@@ -20,6 +20,7 @@
 from pandas.core.dtypes.missing import isna
 
 from pandas.core import algorithms
+from pandas.core.construction import ensure_wrapped_if_datetimelike
 from pandas.core.ops.array_ops import (  # noqa:F401
     arithmetic_op,
     comp_method_OBJECT_ARRAY,
@@ -124,6 +125,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

@@ -174,19 +174,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`.
-    lvalues = ensure_wrapped_if_datetimelike(left)
-    rvalues = ensure_wrapped_if_datetimelike(right)
-    rvalues = _maybe_upcast_for_op(rvalues, lvalues.shape)
+    # NB: We assume that extract_array and ensure_wrapped_if_datetimelike
+    #  has already been called on `left` and `right`.
+    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:
         with np.errstate(all="ignore"):
-            res_values = _na_arithmetic_op(lvalues, rvalues, op)
+            res_values = _na_arithmetic_op(left, right, op)
 
     return res_values
 

pandas/core/series.py

@@ -82,6 +82,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,
@@ -5016,6 +5017,8 @@ def _arith_method(self, other, op):
 
         lvalues = extract_array(self, extract_numpy=True)
         rvalues = extract_array(other, extract_numpy=True)
+        lvalues = ensure_wrapped_if_datetimelike(lvalues)
+        rvalues = ensure_wrapped_if_datetimelike(rvalues)
         result = ops.arithmetic_op(lvalues, rvalues, op)
 
         return self._construct_result(result, name=res_name)