BUG: fix tz-aware DatetimeIndex + TimedeltaIndex (#17558)

doc/source/whatsnew/v0.21.0.txt

@@ -506,6 +506,7 @@ Indexing
 - Bug in ``CategoricalIndex`` reindexing in which specified indices containing duplicates were not being respected (:issue:`17323`)
 - Bug in intersection of ``RangeIndex`` with negative step (:issue:`17296`)
 - Bug in ``IntervalIndex`` where performing a scalar lookup fails for included right endpoints of non-overlapping monotonic decreasing indexes (:issue:`16417`, :issue:`17271`)
+- Bug in adding tz-aware ``DatetimeIndex`` with a numpy array of ``timedelta64``s, and a bug in adding tz-aware ``DatetimeIndex`` with ``TimedeltaIndex`` (:issue:`17558`)
 
 I/O
 ^^^

pandas/core/indexes/datetimelike.py

@@ -13,7 +13,8 @@
     is_integer, is_float,
     is_bool_dtype, _ensure_int64,
     is_scalar, is_dtype_equal,
-    is_list_like)
+    is_timedelta64_dtype, is_datetime64tz_dtype,
+    is_integer_dtype, is_list_like)
 from pandas.core.dtypes.generic import (
     ABCIndex, ABCSeries,
     ABCPeriodIndex, ABCIndexClass)
@@ -651,6 +652,18 @@ def __add__(self, other):
                 raise TypeError("cannot add {typ1} and {typ2}"
                                 .format(typ1=type(self).__name__,
                                         typ2=type(other).__name__))
+            elif isinstance(other, np.ndarray):
+                if is_timedelta64_dtype(other):
+                    return self._add_delta(other)
+                elif is_integer_dtype(other):
+                    # for internal use only:
+                    # allow PeriodIndex + np.array(int64) to
+                    # fallthrough to ufunc operator
+                    return NotImplemented
+                else:
+                    raise TypeError("cannot add {typ1} and np.ndarray[{typ2}]"
+                                    .format(typ1=type(self).__name__,
+                                            typ2=other.dtype))
             elif isinstance(other, (DateOffset, timedelta, np.timedelta64,
                                     Timedelta)):
                 return self._add_delta(other)
@@ -717,7 +730,7 @@ def _add_delta_td(self, other):
 
     def _add_delta_tdi(self, other):
         # add a delta of a TimedeltaIndex
-        # return the i8 result view
+        # return the i8 result view for datetime64tz
 
         # delta operation
         if not len(self) == len(other):
@@ -731,6 +744,8 @@ def _add_delta_tdi(self, other):
         if self.hasnans or other.hasnans:
             mask = (self._isnan) | (other._isnan)
             new_values[mask] = iNaT
+        if is_datetime64tz_dtype(self.dtype):
+            return new_values.view('i8')
         return new_values.view(self.dtype)
 
     def isin(self, values):

pandas/core/indexes/datetimes.py

@@ -14,6 +14,7 @@
     is_integer, is_float,
     is_integer_dtype,
     is_datetime64_ns_dtype,
+    is_timedelta64_dtype,
     is_period_dtype,
     is_bool_dtype,
     is_string_dtype,
@@ -801,6 +802,8 @@ def _add_delta(self, delta):
 
         if isinstance(delta, (Tick, timedelta, np.timedelta64)):
             new_values = self._add_delta_td(delta)
+        elif isinstance(delta, np.ndarray) and is_timedelta64_dtype(delta):
+            new_values = self._add_delta_td(delta)
         elif isinstance(delta, TimedeltaIndex):
             new_values = self._add_delta_tdi(delta)
             # update name when delta is Index

pandas/core/indexes/period.py

@@ -647,7 +647,7 @@ def _maybe_convert_timedelta(self, other):
                 return other.n
             msg = _DIFFERENT_FREQ_INDEX.format(self.freqstr, other.freqstr)
             raise IncompatibleFrequency(msg)
-        elif isinstance(other, np.ndarray):
+        elif isinstance(other, (np.ndarray, TimedeltaIndex)):
             if is_integer_dtype(other):
                 return other
             elif is_timedelta64_dtype(other):

pandas/core/indexes/timedeltas.py

@@ -312,9 +312,11 @@ def _maybe_update_attributes(self, attrs):
         return attrs
 
     def _add_delta(self, delta):
+        name = self.name
         if isinstance(delta, (Tick, timedelta, np.timedelta64)):
             new_values = self._add_delta_td(delta)
-            name = self.name
+        elif isinstance(delta, np.ndarray) and is_timedelta64_dtype(delta):
+            new_values = self._add_delta_td(delta)
         elif isinstance(delta, TimedeltaIndex):
             new_values = self._add_delta_tdi(delta)
             # update name when delta is index

pandas/tests/indexes/datetimelike.py

@@ -1,8 +1,11 @@
 """ generic datetimelike tests """
 
-from .common import Base
+import numpy as np
+import pandas as pd
 import pandas.util.testing as tm
 
+from .common import Base
+
 
 class DatetimeLike(Base):
 
@@ -38,3 +41,30 @@ def test_view(self):
         i_view = i.view(self._holder)
         result = self._holder(i)
         tm.assert_index_equal(result, i_view)
+
+    def test_add_timedelta(self):
+        # GH 17558
+        # Check that tz-aware DatetimeIndex + np.array(dtype="timedelta64")
+        # and DatetimeIndex + TimedeltaIndex work as expected
+        idx = self.create_index()
+        idx.name = "x"
+        if isinstance(idx, pd.DatetimeIndex):
+            idx = idx.tz_localize("US/Eastern")
+
+        expected = idx + np.timedelta64(1, 'D')
+        tm.assert_index_equal(idx, expected - np.timedelta64(1, 'D'))
+
+        deltas = np.array([np.timedelta64(1, 'D')] * len(idx),
+                          dtype="timedelta64[ns]")
+        results = [idx + deltas,     # add numpy array
+                   idx + deltas.astype(dtype="timedelta64[m]"),
+                   idx + pd.TimedeltaIndex(deltas, name=idx.name),
+                   idx + pd.to_timedelta(deltas[0]),
+                   ]
+        for actual in results:
+            tm.assert_index_equal(actual, expected)
+
+        errmsg = (r"cannot add {cls} and np.ndarray\[float64\]"
+                  .format(cls=idx.__class__.__name__))
+        with tm.assert_raises_regex(TypeError, errmsg):
+            idx + np.array([0.1], dtype=np.float64)

pandas/tests/indexes/datetimes/test_datetimelike.py

@@ -76,3 +76,27 @@ def test_union(self):
         for case in cases:
             result = first.union(case)
             assert tm.equalContents(result, everything)
+
+    def test_add_dti_td(self):
+        # GH 17558
+        # Check that tz-aware DatetimeIndex + np.array(dtype="timedelta64")
+        # and DatetimeIndex + TimedeltaIndex work as expected
+        dti = pd.DatetimeIndex([pd.Timestamp("2017/01/01")],
+                               name="x").tz_localize('US/Eastern')
+
+        expected = pd.DatetimeIndex([pd.Timestamp("2017/01/01 01:00")],
+                                    name="x").tz_localize('US/Eastern')
+
+        td_np = np.array([np.timedelta64(1, 'h')], dtype="timedelta64[ns]")
+        results = [dti + td_np,     # add numpy array
+                   dti + td_np.astype(dtype="timedelta64[m]"),
+                   dti + pd.TimedeltaIndex(td_np, name=dti.name),
+                   dti + td_np[0],  # add timedelta scalar
+                   dti + pd.to_timedelta(td_np[0]),
+                   ]
+        for actual in results:
+            tm.assert_index_equal(actual, expected)
+
+        errmsg = r"cannot add DatetimeIndex and np.ndarray\[float64\]"
+        with tm.assert_raises_regex(TypeError, errmsg):
+            dti + np.array([0.1], dtype=np.float64)