Skip to content

Commit c9da215

Browse files
committed
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.
1 parent bee1d27 commit c9da215

File tree

8 files changed

+73
-23
lines changed

8 files changed

+73
-23
lines changed

doc/source/whatsnew/v0.21.0.txt

+1-1
Original file line numberDiff line numberDiff line change
@@ -506,7 +506,7 @@ Indexing
506506
- Bug in ``CategoricalIndex`` reindexing in which specified indices containing duplicates were not being respected (:issue:`17323`)
507507
- Bug in intersection of ``RangeIndex`` with negative step (:issue:`17296`)
508508
- Bug in ``IntervalIndex`` where performing a scalar lookup fails for included right endpoints of non-overlapping monotonic decreasing indexes (:issue:`16417`, :issue:`17271`)
509-
- Bug in adding ``DatetimeIndex`` with a ``TimedeltaIndex`` or a numpy array with ``dtype="timedelta64"`` (:issue:`17558`)
509+
- 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`)
510510

511511
I/O
512512
^^^

pandas/core/indexes/datetimelike.py

+10-5
Original file line numberDiff line numberDiff line change
@@ -13,8 +13,8 @@
1313
is_integer, is_float,
1414
is_bool_dtype, _ensure_int64,
1515
is_scalar, is_dtype_equal,
16-
is_timedelta64_dtype, is_integer_dtype,
17-
is_list_like)
16+
is_timedelta64_dtype, is_datetime64tz_dtype,
17+
is_integer_dtype, is_list_like)
1818
from pandas.core.dtypes.generic import (
1919
ABCIndex, ABCSeries,
2020
ABCPeriodIndex, ABCIndexClass)
@@ -654,8 +654,11 @@ def __add__(self, other):
654654
typ2=type(other).__name__))
655655
elif isinstance(other, np.ndarray):
656656
if is_timedelta64_dtype(other):
657-
return self._add_delta(TimedeltaIndex(other))
657+
return self._add_delta(other)
658658
elif is_integer_dtype(other):
659+
# for internal use only:
660+
# allow PeriodIndex + np.array(int64) to
661+
# fallthrough to ufunc operator
659662
return NotImplemented
660663
else:
661664
raise TypeError("cannot add {typ1} and np.ndarray[{typ2}]"
@@ -727,7 +730,7 @@ def _add_delta_td(self, other):
727730

728731
def _add_delta_tdi(self, other):
729732
# add a delta of a TimedeltaIndex
730-
# return the i8 result view
733+
# return the i8 result view for datetime64tz
731734

732735
# delta operation
733736
if not len(self) == len(other):
@@ -741,7 +744,9 @@ def _add_delta_tdi(self, other):
741744
if self.hasnans or other.hasnans:
742745
mask = (self._isnan) | (other._isnan)
743746
new_values[mask] = iNaT
744-
return new_values.view('i8')
747+
if is_datetime64tz_dtype(self.dtype):
748+
return new_values.view('i8')
749+
return new_values.view(self.dtype)
745750

746751
def isin(self, values):
747752
"""

pandas/core/indexes/datetimes.py

+3
Original file line numberDiff line numberDiff line change
@@ -14,6 +14,7 @@
1414
is_integer, is_float,
1515
is_integer_dtype,
1616
is_datetime64_ns_dtype,
17+
is_timedelta64_dtype,
1718
is_period_dtype,
1819
is_bool_dtype,
1920
is_string_dtype,
@@ -801,6 +802,8 @@ def _add_delta(self, delta):
801802

802803
if isinstance(delta, (Tick, timedelta, np.timedelta64)):
803804
new_values = self._add_delta_td(delta)
805+
elif isinstance(delta, np.ndarray) and is_timedelta64_dtype(delta):
806+
new_values = self._add_delta_td(delta)
804807
elif isinstance(delta, TimedeltaIndex):
805808
new_values = self._add_delta_tdi(delta)
806809
# update name when delta is Index

pandas/core/indexes/period.py

+1-1
Original file line numberDiff line numberDiff line change
@@ -647,7 +647,7 @@ def _maybe_convert_timedelta(self, other):
647647
return other.n
648648
msg = _DIFFERENT_FREQ_INDEX.format(self.freqstr, other.freqstr)
649649
raise IncompatibleFrequency(msg)
650-
elif isinstance(other, np.ndarray):
650+
elif isinstance(other, (np.ndarray, TimedeltaIndex)):
651651
if is_integer_dtype(other):
652652
return other
653653
elif is_timedelta64_dtype(other):

pandas/core/indexes/timedeltas.py

+3-1
Original file line numberDiff line numberDiff line change
@@ -312,9 +312,11 @@ def _maybe_update_attributes(self, attrs):
312312
return attrs
313313

314314
def _add_delta(self, delta):
315+
name = self.name
315316
if isinstance(delta, (Tick, timedelta, np.timedelta64)):
316317
new_values = self._add_delta_td(delta)
317-
name = self.name
318+
elif isinstance(delta, np.ndarray) and is_timedelta64_dtype(delta):
319+
new_values = self._add_delta_td(delta)
318320
elif isinstance(delta, TimedeltaIndex):
319321
new_values = self._add_delta_tdi(delta)
320322
# update name when delta is index

pandas/tests/indexes/datetimelike.py

+31-1
Original file line numberDiff line numberDiff line change
@@ -1,8 +1,10 @@
11
""" generic datetimelike tests """
22

3-
from .common import Base
3+
import numpy as np
4+
import pandas as pd
45
import pandas.util.testing as tm
56

7+
from .common import Base
68

79
class DatetimeLike(Base):
810

@@ -38,3 +40,31 @@ def test_view(self):
3840
i_view = i.view(self._holder)
3941
result = self._holder(i)
4042
tm.assert_index_equal(result, i_view)
43+
44+
def test_add_timedelta(self):
45+
# GH 17558
46+
# Check that tz-aware DatetimeIndex + np.array(dtype="timedelta64")
47+
# and DatetimeIndex + TimedeltaIndex work as expected
48+
idx = self.create_index()
49+
idx.name = "x"
50+
if isinstance(idx, pd.DatetimeIndex):
51+
idx = idx.tz_localize("US/Eastern")
52+
53+
expected = idx + np.timedelta64(1, 'D')
54+
tm.assert_index_equal(idx, expected - np.timedelta64(1, 'D'))
55+
56+
deltas = np.array([np.timedelta64(1, 'D')] * len(idx),
57+
dtype="timedelta64[ns]")
58+
results = [idx + deltas, # add numpy array
59+
idx + deltas.astype(dtype="timedelta64[m]"),
60+
idx + pd.TimedeltaIndex(deltas, name=idx.name),
61+
idx + pd.to_timedelta(deltas[0]),
62+
]
63+
for actual in results:
64+
tm.assert_index_equal(actual, expected)
65+
66+
errmsg = (r"cannot add {cls} and np.ndarray\[float64\]"
67+
.format(cls=idx.__class__.__name__))
68+
with tm.assert_raises_regex(TypeError, errmsg):
69+
idx + np.array([0.1], dtype=np.float64)
70+

pandas/tests/indexes/datetimes/test_datetimelike.py

+24
Original file line numberDiff line numberDiff line change
@@ -76,3 +76,27 @@ def test_union(self):
7676
for case in cases:
7777
result = first.union(case)
7878
assert tm.equalContents(result, everything)
79+
80+
def test_add_dti_td(self):
81+
# GH 17558
82+
# Check that tz-aware DatetimeIndex + np.array(dtype="timedelta64")
83+
# and DatetimeIndex + TimedeltaIndex work as expected
84+
dti = pd.DatetimeIndex([pd.Timestamp("2017/01/01")],
85+
name="x").tz_localize('US/Eastern')
86+
87+
expected = pd.DatetimeIndex([pd.Timestamp("2017/01/01 01:00")],
88+
name="x").tz_localize('US/Eastern')
89+
90+
td_np = np.array([np.timedelta64(1, 'h')], dtype="timedelta64[ns]")
91+
results = [dti + td_np, # add numpy array
92+
dti + td_np.astype(dtype="timedelta64[m]"),
93+
dti + pd.TimedeltaIndex(td_np, name=dti.name),
94+
dti + td_np[0], # add timedelta scalar
95+
dti + pd.to_timedelta(td_np[0]),
96+
]
97+
for actual in results:
98+
tm.assert_index_equal(actual, expected)
99+
100+
errmsg = r"cannot add DatetimeIndex and np.ndarray\[float64\]"
101+
with tm.assert_raises_regex(TypeError, errmsg):
102+
dti + np.array([0.1], dtype=np.float64)

pandas/tests/indexes/datetimes/test_ops.py

-14
Original file line numberDiff line numberDiff line change
@@ -460,20 +460,6 @@ def test_add_dti_dti(self):
460460
with pytest.raises(TypeError):
461461
dti + dti_tz
462462

463-
def test_add_dti_ndarray(self):
464-
# GH 17558
465-
# Check that tz-aware DatetimeIndex + np.array(dtype="timedelta64")
466-
# and DatetimeIndex + TimedeltaIndex work as expected
467-
dti = pd.DatetimeIndex([pd.Timestamp("2017/01/01")])
468-
dti = dti.tz_localize('US/Eastern')
469-
expected = pd.DatetimeIndex([pd.Timestamp("2017/01/01 01:00")])
470-
expected = expected.tz_localize('US/Eastern')
471-
472-
td_np = np.array([np.timedelta64(1, 'h')], dtype="timedelta64[ns]")
473-
tm.assert_index_equal(dti + td_np, expected)
474-
tm.assert_index_equal(dti + td_np[0], expected)
475-
tm.assert_index_equal(dti + TimedeltaIndex(td_np), expected)
476-
477463
def test_difference(self):
478464
for tz in self.tz:
479465
# diff

0 commit comments

Comments
 (0)