Fix/test numeric index ops with timedelta64 array/index/series

pandas-dev · jreback · Aug 23, 2018 · Aug 16, 2018 · Aug 16, 2018 · Aug 16, 2018
commit 6b026911e05e28a1248c4f377b48e0a25e02979d
diff --git a/doc/source/whatsnew/v0.24.0.txt b/doc/source/whatsnew/v0.24.0.txt
@@ -579,6 +579,7 @@ Datetimelike
 - Bug in adding a :class:`Index` with object dtype to a :class:`Series` with ``timedelta64[ns]`` dtype incorrectly raising (:issue:`?????`)
 - Bug in multiplying a :class:`Series` with numeric dtype against a ``timedelta`` object (:issue:`?????`)
 - Bug in :class:`Series` with numeric dtype when adding or subtracting an an array or ``Series`` with ``timedelta64`` dtype (:issue:`?????`)
+- Bug in :class:`Index` with numeric dtype when multiplying or dividing an array with dtype ``timedelta64`` (:issue:`????`)
 
 Timedelta
 ^^^^^^^^^

diff --git a/pandas/core/indexes/base.py b/pandas/core/indexes/base.py
@@ -122,6 +122,14 @@ def index_arithmetic_method(self, other):
         elif isinstance(other, ABCTimedeltaIndex):
             # Defer to subclass implementation
             return NotImplemented
+        elif isinstance(other, np.ndarray) and is_timedelta64_dtype(other):
+            # wrap in Series for op, this will in turn wrap in TimedeltaIndex,
+            # but will correctly raise TypeError instead of NullFrequencyError
+            # for add/sub ops
+            from pandas import Series
+            other = Series(other)
+            out = op(self, other)
+            return Index(out, name=self.name)
 
         other = self._validate_for_numeric_binop(other, op)
 

diff --git a/pandas/core/indexes/range.py b/pandas/core/indexes/range.py
@@ -8,6 +8,7 @@
 from pandas.core.dtypes.common import (
     is_integer,
     is_scalar,
+    is_timedelta64_dtype,
     is_int64_dtype)
 from pandas.core.dtypes.generic import ABCSeries, ABCTimedeltaIndex
 
@@ -596,6 +597,9 @@ def _evaluate_numeric_binop(self, other):
                     # GH#19333 is_integer evaluated True on timedelta64,
                     # so we need to catch these explicitly
                     return op(self._int64index, other)
+                elif is_timedelta64_dtype(other):
+                    # Must be an np.ndarray
+                    return op(self._int64index, other)
 
                 other = self._validate_for_numeric_binop(other, op)
                 attrs = self._get_attributes_dict()

diff --git a/pandas/tests/arithmetic/test_numeric.py b/pandas/tests/arithmetic/test_numeric.py
@@ -47,7 +47,59 @@ def test_operator_series_comparison_zerorank(self):
 # ------------------------------------------------------------------
 # Numeric dtypes Arithmetic with Timedelta Scalar
 
-class TestNumericArraylikeArithmeticWithTimedeltaScalar(object):
+class TestNumericArraylikeArithmeticWithTimedeltaLike(object):
+
+    # TODO: also check name retentention
+    @pytest.mark.parametrize('box_cls', [np.array, pd.Index, pd.Series])
+    @pytest.mark.parametrize('left', [
+        pd.RangeIndex(10, 40, 10)] + [cls([10, 20, 30], dtype=dtype)
+                                      for dtype in ['i1', 'i2', 'i4', 'i8',
+                                                    'u1', 'u2', 'u4', 'u8',
+                                                    'f2', 'f4', 'f8']
+                                      for cls in [pd.Series, pd.Index]
+    ], ids=lambda x: type(x).__name__ + str(x.dtype))
+    def test_mul_td64arr(self, left, box_cls):
+        right = np.array([1, 2, 3], dtype='m8[s]')
+        right = box_cls(right)
+
+        expected = pd.TimedeltaIndex(['10s', '40s', '90s'])
+        if isinstance(left, pd.Series) or box_cls is pd.Series:
+            expected = pd.Series(expected)
+
+        result = left * right
+        tm.assert_equal(result, expected)
+
+        result = right * left
+        tm.assert_equal(result, expected)
+
+    # TODO: also check name retentention
+    @pytest.mark.parametrize('box_cls', [np.array, pd.Index, pd.Series])
+    @pytest.mark.parametrize('left', [
+        pd.RangeIndex(10, 40, 10)] + [cls([10, 20, 30], dtype=dtype)
+                                      for dtype in ['i1', 'i2', 'i4', 'i8',
+                                                    'u1', 'u2', 'u4', 'u8',
+                                                    'f2', 'f4', 'f8']
+                                      for cls in [pd.Series, pd.Index]
+    ], ids=lambda x: type(x).__name__ + str(x.dtype))
+    def test_div_td64arr(self, left, box_cls):
+        right = np.array([10, 40, 90], dtype='m8[s]')
+        right = box_cls(right)
+
+        expected = pd.TimedeltaIndex(['1s', '2s', '3s'])
+        if isinstance(left, pd.Series) or box_cls is pd.Series:
+            expected = pd.Series(expected)
+
+        result = right / left
+        tm.assert_equal(result, expected)
+
+        result = right // left
+        tm.assert_equal(result, expected)
+
+        with pytest.raises(TypeError):
+            left / right
+
+        with pytest.raises(TypeError):
+            left // right
 
     # TODO: de-duplicate with test_numeric_arr_mul_tdscalar
     def test_ops_series(self):