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):