diff --git a/doc/source/whatsnew/v0.21.0.txt b/doc/source/whatsnew/v0.21.0.txt index c808babeee5d9..874657bbbab2f 100644 --- a/doc/source/whatsnew/v0.21.0.txt +++ b/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 ^^^ diff --git a/pandas/core/indexes/datetimelike.py b/pandas/core/indexes/datetimelike.py index c3232627fce74..da3f18ea96dad 100644 --- a/pandas/core/indexes/datetimelike.py +++ b/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): diff --git a/pandas/core/indexes/datetimes.py b/pandas/core/indexes/datetimes.py index 6b1b61c2798f4..b70cfa5f6b75b 100644 --- a/pandas/core/indexes/datetimes.py +++ b/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 diff --git a/pandas/core/indexes/period.py b/pandas/core/indexes/period.py index fb47d1db48610..b5ad94d314610 100644 --- a/pandas/core/indexes/period.py +++ b/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): diff --git a/pandas/core/indexes/timedeltas.py b/pandas/core/indexes/timedeltas.py index d7b7d56d74a3a..5ac822740391b 100644 --- a/pandas/core/indexes/timedeltas.py +++ b/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 diff --git a/pandas/tests/indexes/datetimelike.py b/pandas/tests/indexes/datetimelike.py index 114940009377c..81c5d983987cd 100644 --- a/pandas/tests/indexes/datetimelike.py +++ b/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) diff --git a/pandas/tests/indexes/datetimes/test_datetimelike.py b/pandas/tests/indexes/datetimes/test_datetimelike.py index 538e10e6011ec..ca5b70c2d6f8d 100644 --- a/pandas/tests/indexes/datetimes/test_datetimelike.py +++ b/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)