Fix Timedelta.__floordiv__, __rfloordiv__ (#18961)

Author: jbrockmendel

doc/source/whatsnew/v0.23.0.txt

@@ -449,6 +449,7 @@ Numeric
 - Bug in :func:`Series.__sub__` subtracting a non-nanosecond ``np.datetime64`` object from a ``Series`` gave incorrect results (:issue:`7996`)
 - Bug in :class:`DatetimeIndex`, :class:`TimedeltaIndex` addition and subtraction of zero-dimensional integer arrays gave incorrect results (:issue:`19012`)
 - Bug in :func:`Series.__add__` adding Series with dtype ``timedelta64[ns]`` to a timezone-aware ``DatetimeIndex`` incorrectly dropped timezone information (:issue:`13905`)
+- Bug in :func:`Timedelta.__floordiv__` and :func:`Timedelta.__rfloordiv__` dividing by many incompatible numpy objects was incorrectly allowed (:issue:`18846`)
 -
 
 Categorical

pandas/_libs/tslibs/timedeltas.pyx

@@ -1031,13 +1031,27 @@ class Timedelta(_Timedelta):
         __rdiv__ = __rtruediv__
 
     def __floordiv__(self, other):
+        # numpy does not implement floordiv for timedelta64 dtype, so we cannot
+        # just defer
+        if hasattr(other, '_typ'):
+            # Series, DataFrame, ...
+            return NotImplemented
+
         if hasattr(other, 'dtype'):
-            # work with i8
-            other = other.astype('m8[ns]').astype('i8')
-            return self.value // other
+            if other.dtype.kind == 'm':
+                # also timedelta-like
+                return _broadcast_floordiv_td64(self.value, other, _floordiv)
+            elif other.dtype.kind in ['i', 'u', 'f']:
+                if other.ndim == 0:
+                    return Timedelta(self.value // other)
+                else:
+                    return self.to_timedelta64() // other
+
+            raise TypeError('Invalid dtype {dtype} for '
+                            '{op}'.format(dtype=other.dtype,
+                                          op='__floordiv__'))
 
-        elif is_integer_object(other):
-            # integers only
+        elif is_integer_object(other) or is_float_object(other):
             return Timedelta(self.value // other, unit='ns')
 
         elif not _validate_ops_compat(other):
@@ -1049,20 +1063,79 @@ class Timedelta(_Timedelta):
         return self.value // other.value
 
     def __rfloordiv__(self, other):
-        if hasattr(other, 'dtype'):
-            # work with i8
-            other = other.astype('m8[ns]').astype('i8')
-            return other // self.value
+        # numpy does not implement floordiv for timedelta64 dtype, so we cannot
+        # just defer
+        if hasattr(other, '_typ'):
+            # Series, DataFrame, ...
+            return NotImplemented
 
+        if hasattr(other, 'dtype'):
+            if other.dtype.kind == 'm':
+                # also timedelta-like
+                return _broadcast_floordiv_td64(self.value, other, _rfloordiv)
+            raise TypeError('Invalid dtype {dtype} for '
+                            '{op}'.format(dtype=other.dtype,
+                                          op='__floordiv__'))
+
+        if is_float_object(other) and util._checknull(other):
+            # i.e. np.nan
+            return NotImplemented
         elif not _validate_ops_compat(other):
             return NotImplemented
 
         other = Timedelta(other)
         if other is NaT:
-            return NaT
+            return np.nan
         return other.value // self.value
 
 
+cdef _floordiv(int64_t value, right):
+    return value // right
+
+
+cdef _rfloordiv(int64_t value, right):
+    # analogous to referencing operator.div, but there is no operator.rfloordiv
+    return right // value
+
+
+cdef _broadcast_floordiv_td64(int64_t value, object other,
+                              object (*operation)(int64_t value,
+                                                  object right)):
+    """Boilerplate code shared by Timedelta.__floordiv__ and
+    Timedelta.__rfloordiv__ because np.timedelta64 does not implement these.
+
+    Parameters
+    ----------
+    value : int64_t; `self.value` from a Timedelta object
+    other : object
+    operation : function, either _floordiv or _rfloordiv
+
+    Returns
+    -------
+    result : varies based on `other`
+    """
+    # assumes other.dtype.kind == 'm', i.e. other is timedelta-like
+    cdef:
+        int ndim = getattr(other, 'ndim', -1)
+
+    # We need to watch out for np.timedelta64('NaT').
+    mask = other.view('i8') == NPY_NAT
+
+    if ndim == 0:
+        if mask:
+            return np.nan
+
+        return operation(value, other.astype('m8[ns]').astype('i8'))
+
+    else:
+        res = operation(value, other.astype('m8[ns]').astype('i8'))
+
+        if mask.any():
+            res = res.astype('f8')
+            res[mask] = np.nan
+        return res
+
+
 # resolution in ns
-Timedelta.min = Timedelta(np.iinfo(np.int64).min +1)
+Timedelta.min = Timedelta(np.iinfo(np.int64).min + 1)
 Timedelta.max = Timedelta(np.iinfo(np.int64).max)

pandas/tests/scalar/test_nat.py

@@ -273,6 +273,16 @@ def test_nat_arithmetic():
         assert right - left is NaT
 
 
+def test_nat_rfloordiv_timedelta():
+    # GH#18846
+    # See also test_timedelta.TestTimedeltaArithmetic.test_floordiv
+    td = Timedelta(hours=3, minutes=4)
+
+    assert td // np.nan is NaT
+    assert np.isnan(td // NaT)
+    assert np.isnan(td // np.timedelta64('NaT'))
+
+
 def test_nat_arithmetic_index():
     # GH 11718
 

pandas/tests/scalar/test_timedelta.py

@@ -136,6 +136,7 @@ def test_binary_ops_nat(self):
         assert (td * pd.NaT) is pd.NaT
         assert (td / pd.NaT) is np.nan
         assert (td // pd.NaT) is np.nan
+        assert (td // np.timedelta64('NaT')) is np.nan
 
     def test_binary_ops_integers(self):
         td = Timedelta(10, unit='d')
@@ -162,6 +163,98 @@ def test_binary_ops_with_timedelta(self):
         # invalid multiply with another timedelta
         pytest.raises(TypeError, lambda: td * td)
 
+    def test_floordiv(self):
+        # GH#18846
+        td = Timedelta(hours=3, minutes=4)
+        scalar = Timedelta(hours=3, minutes=3)
+
+        # scalar others
+        assert td // scalar == 1
+        assert -td // scalar.to_pytimedelta() == -2
+        assert (2 * td) // scalar.to_timedelta64() == 2
+
+        assert td // np.nan is pd.NaT
+        assert np.isnan(td // pd.NaT)
+        assert np.isnan(td // np.timedelta64('NaT'))
+
+        with pytest.raises(TypeError):
+            td // np.datetime64('2016-01-01', dtype='datetime64[us]')
+
+        expected = Timedelta(hours=1, minutes=32)
+        assert td // 2 == expected
+        assert td // 2.0 == expected
+        assert td // np.float64(2.0) == expected
+        assert td // np.int32(2.0) == expected
+        assert td // np.uint8(2.0) == expected
+
+        # Array-like others
+        assert td // np.array(scalar.to_timedelta64()) == 1
+
+        res = (3 * td) // np.array([scalar.to_timedelta64()])
+        expected = np.array([3], dtype=np.int64)
+        tm.assert_numpy_array_equal(res, expected)
+
+        res = (10 * td) // np.array([scalar.to_timedelta64(),
+                                     np.timedelta64('NaT')])
+        expected = np.array([10, np.nan])
+        tm.assert_numpy_array_equal(res, expected)
+
+        ser = pd.Series([1], dtype=np.int64)
+        res = td // ser
+        assert res.dtype.kind == 'm'
+
+    def test_rfloordiv(self):
+        # GH#18846
+        td = Timedelta(hours=3, minutes=3)
+        scalar = Timedelta(hours=3, minutes=4)
+
+        # scalar others
+        # x // Timedelta is defined only for timedelta-like x. int-like,
+        # float-like, and date-like, in particular, should all either
+        # a) raise TypeError directly or
+        # b) return NotImplemented, following which the reversed
+        #    operation will raise TypeError.
+        assert td.__rfloordiv__(scalar) == 1
+        assert (-td).__rfloordiv__(scalar.to_pytimedelta()) == -2
+        assert (2 * td).__rfloordiv__(scalar.to_timedelta64()) == 0
+
+        assert np.isnan(td.__rfloordiv__(pd.NaT))
+        assert np.isnan(td.__rfloordiv__(np.timedelta64('NaT')))
+
+        dt64 = np.datetime64('2016-01-01', dtype='datetime64[us]')
+        with pytest.raises(TypeError):
+            td.__rfloordiv__(dt64)
+
+        assert td.__rfloordiv__(np.nan) is NotImplemented
+        assert td.__rfloordiv__(3.5) is NotImplemented
+        assert td.__rfloordiv__(2) is NotImplemented
+
+        with pytest.raises(TypeError):
+            td.__rfloordiv__(np.float64(2.0))
+        with pytest.raises(TypeError):
+            td.__rfloordiv__(np.int32(2.0))
+        with pytest.raises(TypeError):
+            td.__rfloordiv__(np.uint8(9))
+
+        # Array-like others
+        assert td.__rfloordiv__(np.array(scalar.to_timedelta64())) == 1
+
+        res = td.__rfloordiv__(np.array([(3 * scalar).to_timedelta64()]))
+        expected = np.array([3], dtype=np.int64)
+        tm.assert_numpy_array_equal(res, expected)
+
+        arr = np.array([(10 * scalar).to_timedelta64(),
+                        np.timedelta64('NaT')])
+        res = td.__rfloordiv__(arr)
+        expected = np.array([10, np.nan])
+        tm.assert_numpy_array_equal(res, expected)
+
+        ser = pd.Series([1], dtype=np.int64)
+        res = td.__rfloordiv__(ser)
+        assert res is NotImplemented
+        with pytest.raises(TypeError):
+            ser // td
+
 
 class TestTimedeltaComparison(object):
     def test_comparison_object_array(self):