BUG: Fix tz-aware DatetimeIndex +/- TimedeltaIndex/timedelta64 array (#18653)

Author: jschendel

doc/source/whatsnew/v0.22.0.txt

@@ -262,7 +262,7 @@ Indexing
 - Bug in :class:`IntervalIndex` where empty and purely NA data was constructed inconsistently depending on the construction method (:issue:`18421`)
 - Bug in :func:`IntervalIndex.symmetric_difference` where the symmetric difference with a non-``IntervalIndex`` did not raise (:issue:`18475`)
 - Bug in indexing a datetimelike ``Index`` that raised ``ValueError`` instead of ``IndexError`` (:issue:`18386`).
-
+- Bug in tz-aware :class:`DatetimeIndex` where addition/subtraction with a :class:`TimedeltaIndex` or array with ``dtype='timedelta64[ns]'`` was incorrect (:issue:`17558`)
 
 I/O
 ^^^

pandas/core/indexes/datetimelike.py

@@ -14,7 +14,7 @@
     is_integer, is_float,
     is_bool_dtype, _ensure_int64,
     is_scalar, is_dtype_equal,
-    is_list_like)
+    is_list_like, is_timedelta64_dtype)
 from pandas.core.dtypes.generic import (
     ABCIndex, ABCSeries,
     ABCPeriodIndex, ABCIndexClass)
@@ -651,14 +651,14 @@ def __add__(self, other):
             from pandas.core.index import Index
             from pandas.core.indexes.timedeltas import TimedeltaIndex
             from pandas.tseries.offsets import DateOffset
-            if isinstance(other, TimedeltaIndex):
+            if is_timedelta64_dtype(other):
                 return self._add_delta(other)
             elif isinstance(self, TimedeltaIndex) and isinstance(other, Index):
                 if hasattr(other, '_add_delta'):
                     return other._add_delta(self)
                 raise TypeError("cannot add TimedeltaIndex and {typ}"
                                 .format(typ=type(other)))
-            elif isinstance(other, (DateOffset, timedelta, np.timedelta64)):
+            elif isinstance(other, (DateOffset, timedelta)):
                 return self._add_delta(other)
             elif is_integer(other):
                 return self.shift(other)
@@ -674,7 +674,7 @@ def __sub__(self, other):
             from pandas.core.indexes.datetimes import DatetimeIndex
             from pandas.core.indexes.timedeltas import TimedeltaIndex
             from pandas.tseries.offsets import DateOffset
-            if isinstance(other, TimedeltaIndex):
+            if is_timedelta64_dtype(other):
                 return self._add_delta(-other)
             elif isinstance(self, TimedeltaIndex) and isinstance(other, Index):
                 if not isinstance(other, TimedeltaIndex):
@@ -687,7 +687,7 @@ def __sub__(self, other):
                 raise TypeError("cannot subtract {typ1} and {typ2}"
                                 .format(typ1=type(self).__name__,
                                         typ2=type(other).__name__))
-            elif isinstance(other, (DateOffset, timedelta, np.timedelta64)):
+            elif isinstance(other, (DateOffset, timedelta)):
                 return self._add_delta(-other)
             elif is_integer(other):
                 return self.shift(-other)
@@ -736,7 +736,7 @@ def _add_delta_tdi(self, other):
         if self.hasnans or other.hasnans:
             mask = (self._isnan) | (other._isnan)
             new_values[mask] = iNaT
-        return new_values.view(self.dtype)
+        return new_values.view('i8')
 
     def isin(self, values):
         """

pandas/core/indexes/datetimes.py

@@ -13,6 +13,7 @@
     _NS_DTYPE, _INT64_DTYPE,
     is_object_dtype, is_datetime64_dtype,
     is_datetimetz, is_dtype_equal,
+    is_timedelta64_dtype,
     is_integer, is_float,
     is_integer_dtype,
     is_datetime64_ns_dtype,
@@ -858,10 +859,13 @@ def _add_delta(self, delta):
 
         if isinstance(delta, (Tick, timedelta, np.timedelta64)):
             new_values = self._add_delta_td(delta)
-        elif isinstance(delta, TimedeltaIndex):
+        elif is_timedelta64_dtype(delta):
+            if not isinstance(delta, TimedeltaIndex):
+                delta = TimedeltaIndex(delta)
+            else:
+                # update name when delta is Index
+                name = com._maybe_match_name(self, delta)
             new_values = self._add_delta_tdi(delta)
-            # update name when delta is Index
-            name = com._maybe_match_name(self, delta)
         elif isinstance(delta, DateOffset):
             new_values = self._add_offset(delta).asi8
         else:

pandas/tests/indexes/datetimes/test_arithmetic.py

@@ -121,6 +121,99 @@ def test_dti_isub_timedeltalike(self, tz, delta):
         rng -= delta
         tm.assert_index_equal(rng, expected)
 
+    # -------------------------------------------------------------
+    # Binary operations DatetimeIndex and TimedeltaIndex/array
+    def test_dti_add_tdi(self, tz):
+        # GH 17558
+        dti = DatetimeIndex([Timestamp('2017-01-01', tz=tz)] * 10)
+        tdi = pd.timedelta_range('0 days', periods=10)
+        expected = pd.date_range('2017-01-01', periods=10, tz=tz)
+
+        # add with TimdeltaIndex
+        result = dti + tdi
+        tm.assert_index_equal(result, expected)
+
+        result = tdi + dti
+        tm.assert_index_equal(result, expected)
+
+        # add with timedelta64 array
+        result = dti + tdi.values
+        tm.assert_index_equal(result, expected)
+
+        result = tdi.values + dti
+        tm.assert_index_equal(result, expected)
+
+    def test_dti_iadd_tdi(self, tz):
+        # GH 17558
+        dti = DatetimeIndex([Timestamp('2017-01-01', tz=tz)] * 10)
+        tdi = pd.timedelta_range('0 days', periods=10)
+        expected = pd.date_range('2017-01-01', periods=10, tz=tz)
+
+        # iadd with TimdeltaIndex
+        result = DatetimeIndex([Timestamp('2017-01-01', tz=tz)] * 10)
+        result += tdi
+        tm.assert_index_equal(result, expected)
+
+        result = pd.timedelta_range('0 days', periods=10)
+        result += dti
+        tm.assert_index_equal(result, expected)
+
+        # iadd with timedelta64 array
+        result = DatetimeIndex([Timestamp('2017-01-01', tz=tz)] * 10)
+        result += tdi.values
+        tm.assert_index_equal(result, expected)
+
+        result = pd.timedelta_range('0 days', periods=10)
+        result += dti
+        tm.assert_index_equal(result, expected)
+
+    def test_dti_sub_tdi(self, tz):
+        # GH 17558
+        dti = DatetimeIndex([Timestamp('2017-01-01', tz=tz)] * 10)
+        tdi = pd.timedelta_range('0 days', periods=10)
+        expected = pd.date_range('2017-01-01', periods=10, tz=tz, freq='-1D')
+
+        # sub with TimedeltaIndex
+        result = dti - tdi
+        tm.assert_index_equal(result, expected)
+
+        msg = 'cannot subtract TimedeltaIndex and DatetimeIndex'
+        with tm.assert_raises_regex(TypeError, msg):
+            tdi - dti
+
+        # sub with timedelta64 array
+        result = dti - tdi.values
+        tm.assert_index_equal(result, expected)
+
+        msg = 'cannot perform __neg__ with this index type:'
+        with tm.assert_raises_regex(TypeError, msg):
+            tdi.values - dti
+
+    def test_dti_isub_tdi(self, tz):
+        # GH 17558
+        dti = DatetimeIndex([Timestamp('2017-01-01', tz=tz)] * 10)
+        tdi = pd.timedelta_range('0 days', periods=10)
+        expected = pd.date_range('2017-01-01', periods=10, tz=tz, freq='-1D')
+
+        # isub with TimedeltaIndex
+        result = DatetimeIndex([Timestamp('2017-01-01', tz=tz)] * 10)
+        result -= tdi
+        tm.assert_index_equal(result, expected)
+
+        msg = 'cannot subtract TimedeltaIndex and DatetimeIndex'
+        with tm.assert_raises_regex(TypeError, msg):
+            tdi -= dti
+
+        # isub with timedelta64 array
+        result = DatetimeIndex([Timestamp('2017-01-01', tz=tz)] * 10)
+        result -= tdi.values
+        tm.assert_index_equal(result, expected)
+
+        msg = '|'.join(['cannot perform __neg__ with this index type:',
+                        'ufunc subtract cannot use operands with types'])
+        with tm.assert_raises_regex(TypeError, msg):
+            tdi.values -= dti
+
     # -------------------------------------------------------------
     # Binary Operations DatetimeIndex and datetime-like
     # TODO: A couple other tests belong in this section.  Move them in