REF: share TimedeltaArray division code (#50441)

* REF: de-duplicate TimedeltaArray division code

* Share more

* min versions build

Author: jbrockmendel

pandas/core/arrays/timedeltas.py

@@ -1,6 +1,7 @@
 from __future__ import annotations
 
 from datetime import timedelta
+import operator
 from typing import (
     TYPE_CHECKING,
     Iterator,
@@ -65,7 +66,7 @@
 from pandas.core.arrays import datetimelike as dtl
 from pandas.core.arrays._ranges import generate_regular_range
 import pandas.core.common as com
-from pandas.core.construction import extract_array
+from pandas.core.ops import roperator
 from pandas.core.ops.common import unpack_zerodim_and_defer
 
 if TYPE_CHECKING:
@@ -492,10 +493,11 @@ def __mul__(self, other) -> TimedeltaArray:
 
     __rmul__ = __mul__
 
-    @unpack_zerodim_and_defer("__truediv__")
-    def __truediv__(self, other):
-        # timedelta / X is well-defined for timedelta-like or numeric X
-
+    def _scalar_divlike_op(self, other, op):
+        """
+        Shared logic for __truediv__, __rtruediv__, __floordiv__, __rfloordiv__
+        with scalar 'other'.
+        """
         if isinstance(other, self._recognized_scalars):
             other = Timedelta(other)
             # mypy assumes that __new__ returns an instance of the class
@@ -507,31 +509,86 @@ def __truediv__(self, other):
                 return result
 
             # otherwise, dispatch to Timedelta implementation
-            return self._ndarray / other
+            return op(self._ndarray, other)
 
-        elif lib.is_scalar(other):
-            # assume it is numeric
-            result = self._ndarray / other
+        else:
+            # caller is responsible for checking lib.is_scalar(other)
+            # assume other is numeric, otherwise numpy will raise
+
+            if op in [roperator.rtruediv, roperator.rfloordiv]:
+                raise TypeError(
+                    f"Cannot divide {type(other).__name__} by {type(self).__name__}"
+                )
+
+            result = op(self._ndarray, other)
             freq = None
+
             if self.freq is not None:
-                # Tick division is not implemented, so operate on Timedelta
-                freq = self.freq.delta / other
-                freq = to_offset(freq)
+                # Note: freq gets division, not floor-division, even if op
+                #  is floordiv.
+                freq = self.freq / other
+
+                # TODO: 2022-12-24 test_ufunc_coercions, test_tdi_ops_attributes
+                #  get here for truediv, no tests for floordiv
+
+                if op is operator.floordiv:
+                    if freq.nanos == 0 and self.freq.nanos != 0:
+                        # e.g. if self.freq is Nano(1) then dividing by 2
+                        #  rounds down to zero
+                        # TODO: 2022-12-24 should implement the same check
+                        #  for truediv case
+                        freq = None
+
             return type(self)._simple_new(result, dtype=result.dtype, freq=freq)
 
+    def _cast_divlike_op(self, other):
         if not hasattr(other, "dtype"):
             # e.g. list, tuple
             other = np.array(other)
 
         if len(other) != len(self):
             raise ValueError("Cannot divide vectors with unequal lengths")
+        return other
 
-        if is_timedelta64_dtype(other.dtype):
-            # let numpy handle it
-            return self._ndarray / other
+    def _vector_divlike_op(self, other, op) -> np.ndarray | TimedeltaArray:
+        """
+        Shared logic for __truediv__, __floordiv__, and their reversed versions
+        with timedelta64-dtype ndarray other.
+        """
+        # Let numpy handle it
+        result = op(self._ndarray, np.asarray(other))
 
-        elif is_object_dtype(other.dtype):
-            other = extract_array(other, extract_numpy=True)
+        if (is_integer_dtype(other.dtype) or is_float_dtype(other.dtype)) and op in [
+            operator.truediv,
+            operator.floordiv,
+        ]:
+            return type(self)._simple_new(result, dtype=result.dtype)
+
+        if op in [operator.floordiv, roperator.rfloordiv]:
+            mask = self.isna() | isna(other)
+            if mask.any():
+                result = result.astype(np.float64)
+                np.putmask(result, mask, np.nan)
+
+        return result
+
+    @unpack_zerodim_and_defer("__truediv__")
+    def __truediv__(self, other):
+        # timedelta / X is well-defined for timedelta-like or numeric X
+        op = operator.truediv
+        if is_scalar(other):
+            return self._scalar_divlike_op(other, op)
+
+        other = self._cast_divlike_op(other)
+        if (
+            is_timedelta64_dtype(other.dtype)
+            or is_integer_dtype(other.dtype)
+            or is_float_dtype(other.dtype)
+        ):
+            return self._vector_divlike_op(other, op)
+
+        if is_object_dtype(other.dtype):
+            other = np.asarray(other)
             if self.ndim > 1:
                 res_cols = [left / right for left, right in zip(self, other)]
                 res_cols2 = [x.reshape(1, -1) for x in res_cols]
@@ -542,40 +599,18 @@ def __truediv__(self, other):
             return result
 
         else:
-            result = self._ndarray / other
-            return type(self)._simple_new(result, dtype=result.dtype)
+            return NotImplemented
 
     @unpack_zerodim_and_defer("__rtruediv__")
     def __rtruediv__(self, other):
         # X / timedelta is defined only for timedelta-like X
-        if isinstance(other, self._recognized_scalars):
-            other = Timedelta(other)
-            # mypy assumes that __new__ returns an instance of the class
-            # github.com/python/mypy/issues/1020
-            if cast("Timedelta | NaTType", other) is NaT:
-                # specifically timedelta64-NaT
-                result = np.empty(self.shape, dtype=np.float64)
-                result.fill(np.nan)
-                return result
-
-            # otherwise, dispatch to Timedelta implementation
-            return other / self._ndarray
-
-        elif lib.is_scalar(other):
-            raise TypeError(
-                f"Cannot divide {type(other).__name__} by {type(self).__name__}"
-            )
-
-        if not hasattr(other, "dtype"):
-            # e.g. list, tuple
-            other = np.array(other)
-
-        if len(other) != len(self):
-            raise ValueError("Cannot divide vectors with unequal lengths")
+        op = roperator.rtruediv
+        if is_scalar(other):
+            return self._scalar_divlike_op(other, op)
 
+        other = self._cast_divlike_op(other)
         if is_timedelta64_dtype(other.dtype):
-            # let numpy handle it
-            return other / self._ndarray
+            return self._vector_divlike_op(other, op)
 
         elif is_object_dtype(other.dtype):
             # Note: unlike in __truediv__, we do not _need_ to do type
@@ -585,60 +620,24 @@ def __rtruediv__(self, other):
             return np.array(result_list)
 
         else:
-            raise TypeError(
-                f"Cannot divide {other.dtype} data by {type(self).__name__}"
-            )
+            return NotImplemented
 
     @unpack_zerodim_and_defer("__floordiv__")
     def __floordiv__(self, other):
-
+        op = operator.floordiv
         if is_scalar(other):
-            if isinstance(other, self._recognized_scalars):
-                other = Timedelta(other)
-                # mypy assumes that __new__ returns an instance of the class
-                # github.com/python/mypy/issues/1020
-                if cast("Timedelta | NaTType", other) is NaT:
-                    # treat this specifically as timedelta-NaT
-                    result = np.empty(self.shape, dtype=np.float64)
-                    result.fill(np.nan)
-                    return result
-
-                # dispatch to Timedelta implementation
-                return other.__rfloordiv__(self._ndarray)
-
-            # at this point we should only have numeric scalars; anything
-            #  else will raise
-            result = self._ndarray // other
-            freq = None
-            if self.freq is not None:
-                # Note: freq gets division, not floor-division
-                freq = self.freq / other
-                if freq.nanos == 0 and self.freq.nanos != 0:
-                    # e.g. if self.freq is Nano(1) then dividing by 2
-                    #  rounds down to zero
-                    freq = None
-            return type(self)(result, freq=freq)
-
-        if not hasattr(other, "dtype"):
-            # list, tuple
-            other = np.array(other)
-        if len(other) != len(self):
-            raise ValueError("Cannot divide with unequal lengths")
+            return self._scalar_divlike_op(other, op)
 
-        if is_timedelta64_dtype(other.dtype):
-            other = type(self)(other)
-
-            # numpy timedelta64 does not natively support floordiv, so operate
-            #  on the i8 values
-            result = self.asi8 // other.asi8
-            mask = self._isnan | other._isnan
-            if mask.any():
-                result = result.astype(np.float64)
-                np.putmask(result, mask, np.nan)
-            return result
+        other = self._cast_divlike_op(other)
+        if (
+            is_timedelta64_dtype(other.dtype)
+            or is_integer_dtype(other.dtype)
+            or is_float_dtype(other.dtype)
+        ):
+            return self._vector_divlike_op(other, op)
 
         elif is_object_dtype(other.dtype):
-            other = extract_array(other, extract_numpy=True)
+            other = np.asarray(other)
             if self.ndim > 1:
                 res_cols = [left // right for left, right in zip(self, other)]
                 res_cols2 = [x.reshape(1, -1) for x in res_cols]
@@ -649,61 +648,26 @@ def __floordiv__(self, other):
             assert result.dtype == object
             return result
 
-        elif is_integer_dtype(other.dtype) or is_float_dtype(other.dtype):
-            result = self._ndarray // other
-            return type(self)(result)
-
         else:
-            dtype = getattr(other, "dtype", type(other).__name__)
-            raise TypeError(f"Cannot divide {dtype} by {type(self).__name__}")
+            return NotImplemented
 
     @unpack_zerodim_and_defer("__rfloordiv__")
     def __rfloordiv__(self, other):
-
+        op = roperator.rfloordiv
         if is_scalar(other):
-            if isinstance(other, self._recognized_scalars):
-                other = Timedelta(other)
-                # mypy assumes that __new__ returns an instance of the class
-                # github.com/python/mypy/issues/1020
-                if cast("Timedelta | NaTType", other) is NaT:
-                    # treat this specifically as timedelta-NaT
-                    result = np.empty(self.shape, dtype=np.float64)
-                    result.fill(np.nan)
-                    return result
-
-                # dispatch to Timedelta implementation
-                return other.__floordiv__(self._ndarray)
-
-            raise TypeError(
-                f"Cannot divide {type(other).__name__} by {type(self).__name__}"
-            )
-
-        if not hasattr(other, "dtype"):
-            # list, tuple
-            other = np.array(other)
-
-        if len(other) != len(self):
-            raise ValueError("Cannot divide with unequal lengths")
+            return self._scalar_divlike_op(other, op)
 
+        other = self._cast_divlike_op(other)
         if is_timedelta64_dtype(other.dtype):
-            other = type(self)(other)
-            # numpy timedelta64 does not natively support floordiv, so operate
-            #  on the i8 values
-            result = other.asi8 // self.asi8
-            mask = self._isnan | other._isnan
-            if mask.any():
-                result = result.astype(np.float64)
-                np.putmask(result, mask, np.nan)
-            return result
+            return self._vector_divlike_op(other, op)
 
         elif is_object_dtype(other.dtype):
             result_list = [other[n] // self[n] for n in range(len(self))]
             result = np.array(result_list)
             return result
 
         else:
-            dtype = getattr(other, "dtype", type(other).__name__)
-            raise TypeError(f"Cannot divide {dtype} by {type(self).__name__}")
+            return NotImplemented
 
     @unpack_zerodim_and_defer("__mod__")
     def __mod__(self, other):

pandas/tests/arithmetic/test_numeric.py

@@ -177,10 +177,12 @@ def test_div_td64arr(self, left, box_cls):
         result = right // left
         tm.assert_equal(result, expected)
 
-        msg = "Cannot divide"
+        # (true_) needed for min-versions build 2022-12-26
+        msg = "ufunc '(true_)?divide' cannot use operands with types"
         with pytest.raises(TypeError, match=msg):
             left / right
 
+        msg = "ufunc 'floor_divide' cannot use operands with types"
         with pytest.raises(TypeError, match=msg):
             left // right
 

pandas/tests/arithmetic/test_timedelta64.py

@@ -2019,6 +2019,7 @@ def test_td64arr_div_numeric_array(
                 "cannot perform __truediv__",
                 "unsupported operand",
                 "Cannot divide",
+                "ufunc 'divide' cannot use operands with types",
             ]
         )
         with pytest.raises(TypeError, match=pattern):