Add DatetimeIndex and timedeltas via _add_delta_td

The internal implementation of _add_delta_td already correctly
adds a datetime index and an array of timedeltas, so use that
implementation instead of wrapping in a TimedeltaIndex (which
requires a bit more metadata like name to be passed around).

Move test to datetimelike to check addition of TimedeltaIndex
and np.array(timedelta64) with each of {DatetimeIndex,
TimedeltaIndex, PeriodIndex}. Fix the latter two to explicitly
support addition with a numpy array.

Clarify some comments such as the whatsnew and move tests to
test_datetimelike.

Author: azjps

doc/source/whatsnew/v0.21.0.txt

@@ -506,7 +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 ``DatetimeIndex`` with a ``TimedeltaIndex`` or a numpy array with ``dtype="timedelta64"`` (:issue:`17558`)
+- 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,8 +13,8 @@
     is_integer, is_float,
     is_bool_dtype, _ensure_int64,
     is_scalar, is_dtype_equal,
-    is_timedelta64_dtype, is_integer_dtype,
-    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)
@@ -654,8 +654,11 @@ def __add__(self, other):
                                         typ2=type(other).__name__))
             elif isinstance(other, np.ndarray):
                 if is_timedelta64_dtype(other):
-                    return self._add_delta(TimedeltaIndex(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}]"
@@ -727,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):
@@ -741,7 +744,9 @@ 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('i8')
+        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,10 @@
 """ 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 +40,31 @@ 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)

pandas/tests/indexes/datetimes/test_ops.py

@@ -460,20 +460,6 @@ def test_add_dti_dti(self):
         with pytest.raises(TypeError):
             dti + dti_tz
 
-    def test_add_dti_ndarray(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")])
-        dti = dti.tz_localize('US/Eastern')
-        expected = pd.DatetimeIndex([pd.Timestamp("2017/01/01 01:00")])
-        expected = expected.tz_localize('US/Eastern')
-
-        td_np = np.array([np.timedelta64(1, 'h')], dtype="timedelta64[ns]")
-        tm.assert_index_equal(dti + td_np, expected)
-        tm.assert_index_equal(dti + td_np[0], expected)
-        tm.assert_index_equal(dti + TimedeltaIndex(td_np), expected)
-
     def test_difference(self):
         for tz in self.tz:
             # diff