REF: simplify Timedelta arithmetic methods (pandas-dev#33978)

Author: jbrockmendel

pandas/_libs/tslibs/timedeltas.pyx

@@ -18,7 +18,8 @@ PyDateTime_IMPORT
 cimport pandas._libs.tslibs.util as util
 from pandas._libs.tslibs.util cimport (
     is_timedelta64_object, is_datetime64_object, is_integer_object,
-    is_float_object)
+    is_float_object, is_array
+)
 
 from pandas._libs.tslibs.c_timestamp cimport _Timestamp
 
@@ -606,7 +607,7 @@ def _binary_op_method_timedeltalike(op, name):
             # We are implicitly requiring the canonical behavior to be
             # defined by Timestamp methods.
 
-        elif hasattr(other, 'dtype'):
+        elif is_array(other):
             # nd-array like
             if other.dtype.kind in ['m', 'M']:
                 return op(self.to_timedelta64(), other)
@@ -1347,113 +1348,64 @@ class Timedelta(_Timedelta):
     __rsub__ = _binary_op_method_timedeltalike(lambda x, y: y - x, '__rsub__')
 
     def __mul__(self, other):
-        if hasattr(other, '_typ'):
-            # Series, DataFrame, ...
-            if other._typ == 'dateoffset' and hasattr(other, 'delta'):
-                # Tick offset; this op will raise TypeError
-                return other.delta * self
-            return NotImplemented
+        if is_integer_object(other) or is_float_object(other):
+            return Timedelta(other * self.value, unit='ns')
 
-        elif util.is_nan(other):
-            # i.e. np.nan, but also catch np.float64("NaN") which would
-            #  otherwise get caught by the hasattr(other, "dtype") branch
-            #  incorrectly return a np.timedelta64 object.
-            return NaT
-
-        elif hasattr(other, 'dtype'):
+        elif is_array(other):
             # ndarray-like
             return other * self.to_timedelta64()
 
-        elif other is NaT:
-            raise TypeError('Cannot multiply Timedelta with NaT')
-
-        elif not (is_integer_object(other) or is_float_object(other)):
-            # only integers and floats allowed
-            return NotImplemented
-
-        return Timedelta(other * self.value, unit='ns')
+        return NotImplemented
 
     __rmul__ = __mul__
 
     def __truediv__(self, other):
-        if hasattr(other, '_typ'):
-            # Series, DataFrame, ...
-            if other._typ == 'dateoffset' and hasattr(other, 'delta'):
-                # Tick offset
-                return self / other.delta
-            return NotImplemented
-
-        elif is_timedelta64_object(other):
-            # convert to Timedelta below
-            pass
-
-        elif util.is_nan(other):
-            # i.e. np.nan, but also catch np.float64("NaN") which would
-            #  otherwise get caught by the hasattr(other, "dtype") branch
-            #  incorrectly return a np.timedelta64 object.
-            return NaT
-
-        elif hasattr(other, 'dtype'):
-            return self.to_timedelta64() / other
+        if _should_cast_to_timedelta(other):
+            # We interpret NaT as timedelta64("NaT")
+            other = Timedelta(other)
+            if other is NaT:
+                return np.nan
+            return self.value / float(other.value)
 
         elif is_integer_object(other) or is_float_object(other):
             # integers or floats
             return Timedelta(self.value / other, unit='ns')
 
-        elif not _validate_ops_compat(other):
-            return NotImplemented
+        elif is_array(other):
+            return self.to_timedelta64() / other
 
-        other = Timedelta(other)
-        if other is NaT:
-            return np.nan
-        return self.value / float(other.value)
+        return NotImplemented
 
     def __rtruediv__(self, other):
-        if hasattr(other, '_typ'):
-            # Series, DataFrame, ...
-            if other._typ == 'dateoffset' and hasattr(other, 'delta'):
-                # Tick offset
-                return other.delta / self
-            return NotImplemented
-
-        elif is_timedelta64_object(other):
-            # convert to Timedelta below
-            pass
-
-        elif util.is_nan(other):
-            # i.e. np.nan or np.float64("NaN")
-            raise TypeError("Cannot divide float by Timedelta")
+        if _should_cast_to_timedelta(other):
+            # We interpret NaT as timedelta64("NaT")
+            other = Timedelta(other)
+            if other is NaT:
+                return np.nan
+            return float(other.value) / self.value
 
-        elif hasattr(other, 'dtype'):
+        elif is_array(other):
             if other.dtype.kind == "O":
                 # GH#31869
                 return np.array([x / self for x in other])
             return other / self.to_timedelta64()
 
-        elif not _validate_ops_compat(other):
-            return NotImplemented
-
-        other = Timedelta(other)
-        if other is NaT:
-            # In this context we treat NaT as timedelta-like
-            return np.nan
-        return float(other.value) / self.value
+        return NotImplemented
 
     def __floordiv__(self, other):
         # numpy does not implement floordiv for timedelta64 dtype, so we cannot
         # just defer
-        if hasattr(other, '_typ'):
-            # Series, DataFrame, ...
-            if other._typ == 'dateoffset' and hasattr(other, 'delta'):
-                # Tick offset
-                return self // other.delta
-            return NotImplemented
+        if _should_cast_to_timedelta(other):
+            # We interpret NaT as timedelta64("NaT")
+            other = Timedelta(other)
+            if other is NaT:
+                return np.nan
+            return self.value // other.value
 
-        elif is_timedelta64_object(other):
-            # convert to Timedelta below
-            pass
+        elif is_integer_object(other) or is_float_object(other):
+            return Timedelta(self.value // other, unit='ns')
 
-        elif hasattr(other, 'dtype'):
+        elif is_array(other):
             if other.dtype.kind == 'm':
                 # also timedelta-like
                 return _broadcast_floordiv_td64(self.value, other, _floordiv)
@@ -1465,50 +1417,27 @@ class Timedelta(_Timedelta):
 
             raise TypeError(f'Invalid dtype {other.dtype} for __floordiv__')
 
-        elif is_integer_object(other) or is_float_object(other):
-            return Timedelta(self.value // other, unit='ns')
-
-        elif not _validate_ops_compat(other):
-            return NotImplemented
-
-        other = Timedelta(other)
-        if other is NaT:
-            return np.nan
-        return self.value // other.value
+        return NotImplemented
 
     def __rfloordiv__(self, other):
         # numpy does not implement floordiv for timedelta64 dtype, so we cannot
         # just defer
-        if hasattr(other, '_typ'):
-            # Series, DataFrame, ...
-            if other._typ == 'dateoffset' and hasattr(other, 'delta'):
-                # Tick offset
-                return other.delta // self
-            return NotImplemented
-
-        elif is_timedelta64_object(other):
-            # convert to Timedelta below
-            pass
+        if _should_cast_to_timedelta(other):
+            # We interpret NaT as timedelta64("NaT")
+            other = Timedelta(other)
+            if other is NaT:
+                return np.nan
+            return other.value // self.value
 
-        elif hasattr(other, 'dtype'):
+        elif is_array(other):
             if other.dtype.kind == 'm':
                 # also timedelta-like
                 return _broadcast_floordiv_td64(self.value, other, _rfloordiv)
 
             # Includes integer array // Timedelta, disallowed in GH#19761
             raise TypeError(f'Invalid dtype {other.dtype} for __floordiv__')
 
-        elif is_float_object(other) and util.is_nan(other):
-            # i.e. np.nan
-            return NotImplemented
-
-        elif not _validate_ops_compat(other):
-            return NotImplemented
-
-        other = Timedelta(other)
-        if other is NaT:
-            return np.nan
-        return other.value // self.value
+        return NotImplemented
 
     def __mod__(self, other):
         # Naive implementation, room for optimization
@@ -1529,6 +1458,21 @@ class Timedelta(_Timedelta):
         return div, other - div * self
 
 
+cdef bint is_any_td_scalar(object obj):
+    return (
+        PyDelta_Check(obj) or is_timedelta64_object(obj) or isinstance(obj, Tick)
+    )
+
+
+cdef bint _should_cast_to_timedelta(object obj):
+    """
+    Should we treat this object as a Timedelta for the purpose of a binary op
+    """
+    return (
+        is_any_td_scalar(obj) or obj is None or obj is NaT or isinstance(obj, str)
+    )
+
+
 cdef _floordiv(int64_t value, right):
     return value // right
 

pandas/tests/scalar/test_nat.py

@@ -389,7 +389,8 @@ def test_nat_arithmetic_scalar(op_name, value, val_type):
             and "times" in op_name
             and isinstance(value, Timedelta)
         ):
-            msg = "Cannot multiply"
+            typs = "(Timedelta|NaTType)"
+            msg = rf"unsupported operand type\(s\) for \*: '{typs}' and '{typs}'"
         elif val_type == "str":
             # un-specific check here because the message comes from str
             #  and varies by method

pandas/tests/scalar/timedelta/test_arithmetic.py

@@ -325,7 +325,13 @@ class TestTimedeltaMultiplicationDivision:
     def test_td_mul_nat(self, op, td_nat):
         # GH#19819
         td = Timedelta(10, unit="d")
-        msg = "cannot use operands with types|Cannot multiply Timedelta with NaT"
+        typs = "|".join(["numpy.timedelta64", "NaTType", "Timedelta"])
+        msg = "|".join(
+            [
+                rf"unsupported operand type\(s\) for \*: '{typs}' and '{typs}'",
+                r"ufunc '?multiply'? cannot use operands with types",
+            ]
+        )
         with pytest.raises(TypeError, match=msg):
             op(td, td_nat)
 
@@ -457,11 +463,11 @@ def test_td_rdiv_na_scalar(self):
         result = np.timedelta64("NaT") / td
         assert np.isnan(result)
 
-        msg = "cannot use operands with types dtype"
+        msg = r"unsupported operand type\(s\) for /: 'numpy.datetime64' and 'Timedelta'"
         with pytest.raises(TypeError, match=msg):
             np.datetime64("NaT") / td
 
-        msg = "Cannot divide float by Timedelta"
+        msg = r"unsupported operand type\(s\) for /: 'float' and 'Timedelta'"
         with pytest.raises(TypeError, match=msg):
             np.nan / td
 
@@ -479,7 +485,7 @@ def test_td_rdiv_ndarray(self):
         tm.assert_numpy_array_equal(result, expected)
 
         arr = np.array([np.nan], dtype=object)
-        msg = "Cannot divide float by Timedelta"
+        msg = r"unsupported operand type\(s\) for /: 'float' and 'Timedelta'"
         with pytest.raises(TypeError, match=msg):
             arr / td
 
@@ -522,6 +528,7 @@ def test_td_floordiv_invalid_scalar(self):
             [
                 r"Invalid dtype datetime64\[D\] for __floordiv__",
                 "'dtype' is an invalid keyword argument for this function",
+                r"ufunc '?floor_divide'? cannot use operands with types",
             ]
         )
         with pytest.raises(TypeError, match=msg):
@@ -595,9 +602,14 @@ def test_td_rfloordiv_invalid_scalar(self):
         td = Timedelta(hours=3, minutes=3)
 
         dt64 = np.datetime64("2016-01-01", "us")
-        msg = r"Invalid dtype datetime64\[us\] for __floordiv__"
+
+        assert td.__rfloordiv__(dt64) is NotImplemented
+
+        msg = (
+            r"unsupported operand type\(s\) for //: 'numpy.datetime64' and 'Timedelta'"
+        )
         with pytest.raises(TypeError, match=msg):
-            td.__rfloordiv__(dt64)
+            dt64 // td
 
     def test_td_rfloordiv_numeric_scalar(self):
         # GH#18846
@@ -606,15 +618,18 @@ def test_td_rfloordiv_numeric_scalar(self):
         assert td.__rfloordiv__(np.nan) is NotImplemented
         assert td.__rfloordiv__(3.5) is NotImplemented
         assert td.__rfloordiv__(2) is NotImplemented
+        assert td.__rfloordiv__(np.float64(2.0)) is NotImplemented
+        assert td.__rfloordiv__(np.uint8(9)) is NotImplemented
+        assert td.__rfloordiv__(np.int32(2.0)) is NotImplemented
 
-        msg = "Invalid dtype"
+        msg = r"unsupported operand type\(s\) for //: '.*' and 'Timedelta"
         with pytest.raises(TypeError, match=msg):
-            td.__rfloordiv__(np.float64(2.0))
+            np.float64(2.0) // td
         with pytest.raises(TypeError, match=msg):
-            td.__rfloordiv__(np.uint8(9))
+            np.uint8(9) // td
         with pytest.raises(TypeError, match=msg):
             # deprecated GH#19761, enforced GH#29797
-            td.__rfloordiv__(np.int32(2.0))
+            np.int32(2.0) // td
 
     def test_td_rfloordiv_timedeltalike_array(self):
         # GH#18846