Fix+test DTI/TDI/PI add/sub with ndarray[datetime64/timedelta64] (#19847)

Author: jbrockmendel

doc/source/whatsnew/v0.23.0.txt

@@ -756,6 +756,7 @@ Datetimelike
 - Bug in :func:`Timestamp.floor` :func:`DatetimeIndex.floor` where time stamps far in the future and past were not rounded correctly (:issue:`19206`)
 - Bug in :func:`to_datetime` where passing an out-of-bounds datetime with ``errors='coerce'`` and ``utc=True`` would raise ``OutOfBoundsDatetime`` instead of parsing to ``NaT`` (:issue:`19612`)
 - Bug in :class:`DatetimeIndex` and :class:`TimedeltaIndex` addition and subtraction where name of the returned object was not always set consistently. (:issue:`19744`)
+- Bug in :class:`DatetimeIndex` and :class:`TimedeltaIndex` addition and subtraction where operations with numpy arrays raised ``TypeError`` (:issue:`19847`)
 -
 
 Timedelta

pandas/core/indexes/datetimelike.py

@@ -31,6 +31,7 @@
     is_integer_dtype,
     is_object_dtype,
     is_string_dtype,
+    is_datetime64_dtype,
     is_period_dtype,
     is_timedelta64_dtype)
 from pandas.core.dtypes.generic import (
@@ -676,9 +677,7 @@ def _add_datetimelike_methods(cls):
         """
 
         def __add__(self, other):
-            from pandas.core.index import Index
-            from pandas.core.indexes.timedeltas import TimedeltaIndex
-            from pandas.tseries.offsets import DateOffset
+            from pandas import Index, DatetimeIndex, TimedeltaIndex, DateOffset
 
             other = lib.item_from_zerodim(other)
             if isinstance(other, ABCSeries):
@@ -710,6 +709,9 @@ def __add__(self, other):
                                     .format(typ=type(other)))
             elif isinstance(other, Index):
                 result = self._add_datelike(other)
+            elif is_datetime64_dtype(other):
+                # ndarray[datetime64]; note DatetimeIndex is caught above
+                return self + DatetimeIndex(other)
             elif is_integer_dtype(other) and self.freq is None:
                 # GH#19123
                 raise NullFrequencyError("Cannot shift with no freq")
@@ -729,10 +731,7 @@ def __radd__(self, other):
         cls.__radd__ = __radd__
 
         def __sub__(self, other):
-            from pandas.core.index import Index
-            from pandas.core.indexes.datetimes import DatetimeIndex
-            from pandas.core.indexes.timedeltas import TimedeltaIndex
-            from pandas.tseries.offsets import DateOffset
+            from pandas import Index, DatetimeIndex, TimedeltaIndex, DateOffset
 
             other = lib.item_from_zerodim(other)
             if isinstance(other, ABCSeries):
@@ -764,6 +763,9 @@ def __sub__(self, other):
                                 .format(typ=type(other).__name__))
             elif isinstance(other, DatetimeIndex):
                 result = self._sub_datelike(other)
+            elif is_datetime64_dtype(other):
+                # ndarray[datetime64]; note we caught DatetimeIndex earlier
+                return self - DatetimeIndex(other)
             elif isinstance(other, Index):
                 raise TypeError("cannot subtract {typ1} and {typ2}"
                                 .format(typ1=type(self).__name__,
@@ -782,6 +784,11 @@ def __sub__(self, other):
         cls.__sub__ = __sub__
 
         def __rsub__(self, other):
+            if is_datetime64_dtype(other) and is_timedelta64_dtype(self):
+                # ndarray[datetime64] cannot be subtracted from self, so
+                # we need to wrap in DatetimeIndex and flip the operation
+                from pandas import DatetimeIndex
+                return DatetimeIndex(other) - self
             return -(self - other)
         cls.__rsub__ = __rsub__
 

pandas/core/indexes/timedeltas.py

@@ -377,6 +377,10 @@ def _add_delta(self, delta):
             new_values = self._add_delta_td(delta)
         elif isinstance(delta, TimedeltaIndex):
             new_values = self._add_delta_tdi(delta)
+        elif is_timedelta64_dtype(delta):
+            # ndarray[timedelta64] --> wrap in TimedeltaIndex
+            delta = TimedeltaIndex(delta)
+            new_values = self._add_delta_tdi(delta)
         else:
             raise TypeError("cannot add the type {0} to a TimedeltaIndex"
                             .format(type(delta)))

pandas/tests/indexes/datetimes/test_arithmetic.py

@@ -571,6 +571,62 @@ def test_add_datetimelike_and_dti_tz(self, addend):
         with tm.assert_raises_regex(TypeError, msg):
             addend + dti_tz
 
+    # -------------------------------------------------------------
+    # __add__/__sub__ with ndarray[datetime64] and ndarray[timedelta64]
+
+    def test_dti_add_dt64_array_raises(self, tz):
+        dti = pd.date_range('2016-01-01', periods=3, tz=tz)
+        dtarr = dti.values
+
+        with pytest.raises(TypeError):
+            dti + dtarr
+        with pytest.raises(TypeError):
+            dtarr + dti
+
+    def test_dti_sub_dt64_array_naive(self):
+        dti = pd.date_range('2016-01-01', periods=3, tz=None)
+        dtarr = dti.values
+
+        expected = dti - dti
+        result = dti - dtarr
+        tm.assert_index_equal(result, expected)
+        result = dtarr - dti
+        tm.assert_index_equal(result, expected)
+
+    def test_dti_sub_dt64_array_aware_raises(self, tz):
+        if tz is None:
+            return
+        dti = pd.date_range('2016-01-01', periods=3, tz=tz)
+        dtarr = dti.values
+
+        with pytest.raises(TypeError):
+            dti - dtarr
+        with pytest.raises(TypeError):
+            dtarr - dti
+
+    def test_dti_add_td64_array(self, tz):
+        dti = pd.date_range('2016-01-01', periods=3, tz=tz)
+        tdi = dti - dti.shift(1)
+        tdarr = tdi.values
+
+        expected = dti + tdi
+        result = dti + tdarr
+        tm.assert_index_equal(result, expected)
+        result = tdarr + dti
+        tm.assert_index_equal(result, expected)
+
+    def test_dti_sub_td64_array(self, tz):
+        dti = pd.date_range('2016-01-01', periods=3, tz=tz)
+        tdi = dti - dti.shift(1)
+        tdarr = tdi.values
+
+        expected = dti - tdi
+        result = dti - tdarr
+        tm.assert_index_equal(result, expected)
+
+        with pytest.raises(TypeError):
+            tdarr - dti
+
     # -------------------------------------------------------------
 
     def test_sub_dti_dti(self):

pandas/tests/indexes/period/test_arithmetic.py

@@ -255,6 +255,64 @@ def test_comp_nat(self, dtype):
 
 
 class TestPeriodIndexArithmetic(object):
+
+    # -----------------------------------------------------------------
+    # __add__/__sub__ with ndarray[datetime64] and ndarray[timedelta64]
+
+    def test_pi_add_sub_dt64_array_raises(self):
+        rng = pd.period_range('1/1/2000', freq='D', periods=3)
+        dti = pd.date_range('2016-01-01', periods=3)
+        dtarr = dti.values
+
+        with pytest.raises(TypeError):
+            rng + dtarr
+        with pytest.raises(TypeError):
+            dtarr + rng
+
+        with pytest.raises(TypeError):
+            rng - dtarr
+        with pytest.raises(TypeError):
+            dtarr - rng
+
+    def test_pi_add_sub_td64_array_non_tick_raises(self):
+        rng = pd.period_range('1/1/2000', freq='Q', periods=3)
+        dti = pd.date_range('2016-01-01', periods=3)
+        tdi = dti - dti.shift(1)
+        tdarr = tdi.values
+
+        with pytest.raises(period.IncompatibleFrequency):
+            rng + tdarr
+        with pytest.raises(period.IncompatibleFrequency):
+            tdarr + rng
+
+        with pytest.raises(period.IncompatibleFrequency):
+            rng - tdarr
+        with pytest.raises(period.IncompatibleFrequency):
+            tdarr - rng
+
+    @pytest.mark.xfail(reason='op with TimedeltaIndex raises, with ndarray OK')
+    def test_pi_add_sub_td64_array_tick(self):
+        rng = pd.period_range('1/1/2000', freq='Q', periods=3)
+        dti = pd.date_range('2016-01-01', periods=3)
+        tdi = dti - dti.shift(1)
+        tdarr = tdi.values
+
+        expected = rng + tdi
+        result = rng + tdarr
+        tm.assert_index_equal(result, expected)
+        result = tdarr + rng
+        tm.assert_index_equal(result, expected)
+
+        expected = rng - tdi
+        result = rng - tdarr
+        tm.assert_index_equal(result, expected)
+
+        with pytest.raises(TypeError):
+            tdarr - rng
+
+    # -----------------------------------------------------------------
+    # operations with array/Index of DateOffset objects
+
     @pytest.mark.parametrize('box', [np.array, pd.Index])
     def test_pi_add_offset_array(self, box):
         # GH#18849

pandas/tests/indexes/timedeltas/test_arithmetic.py

@@ -586,6 +586,55 @@ def test_tdi_radd_timestamp(self):
         expected = DatetimeIndex(['2011-01-02', '2011-01-03'])
         tm.assert_index_equal(result, expected)
 
+    # -------------------------------------------------------------
+    # __add__/__sub__ with ndarray[datetime64] and ndarray[timedelta64]
+
+    def test_tdi_sub_dt64_array(self):
+        dti = pd.date_range('2016-01-01', periods=3)
+        tdi = dti - dti.shift(1)
+        dtarr = dti.values
+
+        with pytest.raises(TypeError):
+            tdi - dtarr
+
+        # TimedeltaIndex.__rsub__
+        expected = pd.DatetimeIndex(dtarr) - tdi
+        result = dtarr - tdi
+        tm.assert_index_equal(result, expected)
+
+    def test_tdi_add_dt64_array(self):
+        dti = pd.date_range('2016-01-01', periods=3)
+        tdi = dti - dti.shift(1)
+        dtarr = dti.values
+
+        expected = pd.DatetimeIndex(dtarr) + tdi
+        result = tdi + dtarr
+        tm.assert_index_equal(result, expected)
+        result = dtarr + tdi
+        tm.assert_index_equal(result, expected)
+
+    def test_tdi_add_td64_array(self):
+        dti = pd.date_range('2016-01-01', periods=3)
+        tdi = dti - dti.shift(1)
+        tdarr = tdi.values
+
+        expected = 2 * tdi
+        result = tdi + tdarr
+        tm.assert_index_equal(result, expected)
+        result = tdarr + tdi
+        tm.assert_index_equal(result, expected)
+
+    def test_tdi_sub_td64_array(self):
+        dti = pd.date_range('2016-01-01', periods=3)
+        tdi = dti - dti.shift(1)
+        tdarr = tdi.values
+
+        expected = 0 * tdi
+        result = tdi - tdarr
+        tm.assert_index_equal(result, expected)
+        result = tdarr - tdi
+        tm.assert_index_equal(result, expected)
+
     # -------------------------------------------------------------
 
     def test_subtraction_ops(self):