Skip to content

Fix+test DTI/TDI/PI add/sub with ndarray[datetime64/timedelta64] #19847

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
merged 7 commits into from
Feb 25, 2018
13 changes: 13 additions & 0 deletions pandas/core/indexes/datetimelike.py
Original file line number Diff line number Diff line change
Expand Up @@ -25,6 +25,7 @@
is_integer_dtype,
is_object_dtype,
is_string_dtype,
is_datetime64_dtype,
is_timedelta64_dtype)
from pandas.core.dtypes.generic import (
ABCIndex, ABCSeries, ABCPeriodIndex, ABCIndexClass)
Expand Down Expand Up @@ -654,6 +655,7 @@ def _add_datetimelike_methods(cls):

def __add__(self, other):
from pandas.core.index import Index
from pandas.core.indexes.datetimes import DatetimeIndex
from pandas.core.indexes.timedeltas import TimedeltaIndex
from pandas.tseries.offsets import DateOffset

Expand Down Expand Up @@ -681,6 +683,9 @@ def __add__(self, other):
result = self._add_datelike(other)
elif isinstance(other, Index):
result = self._add_datelike(other)
elif is_datetime64_dtype(other):
# ndarray[datetime64]; note DatetimeIndex is caught above
return self + DatetimeIndex(other)
elif is_integer_dtype(other) and self.freq is None:
# GH#19123
raise NullFrequencyError("Cannot shift with no freq")
Expand Down Expand Up @@ -726,6 +731,9 @@ def __sub__(self, other):
result = self._sub_datelike(other)
elif isinstance(other, Period):
result = self._sub_period(other)
elif is_datetime64_dtype(other):
# ndarray[datetime64]; note we caught DatetimeIndex earlier
return self - DatetimeIndex(other)
elif isinstance(other, Index):
raise TypeError("cannot subtract {typ1} and {typ2}"
.format(typ1=type(self).__name__,
Expand All @@ -744,6 +752,11 @@ def __sub__(self, other):
cls.__sub__ = __sub__

def __rsub__(self, other):
if is_datetime64_dtype(other) and is_timedelta64_dtype(self):
# ndarray[datetime64] cannot be subtracted from self, so
# we need to wrap in DatetimeIndex and flip the operation
from pandas.core.indexes.datetimes import DatetimeIndex
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

since inside a function, can just import from pandas directrly

Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

atually pls change all imports that are inside functions to do this (its just shorter)

return DatetimeIndex(other) - self
return -(self - other)
cls.__rsub__ = __rsub__

Expand Down
4 changes: 4 additions & 0 deletions pandas/core/indexes/timedeltas.py
Original file line number Diff line number Diff line change
Expand Up @@ -377,6 +377,10 @@ def _add_delta(self, delta):
new_values = self._add_delta_td(delta)
elif isinstance(delta, TimedeltaIndex):
new_values = self._add_delta_tdi(delta)
elif is_timedelta64_dtype(delta):
# ndarray[timedelta64] --> wrap in TimedeltaIndex
delta = TimedeltaIndex(delta)
new_values = self._add_delta_tdi(delta)
else:
raise TypeError("cannot add the type {0} to a TimedeltaIndex"
.format(type(delta)))
Expand Down
56 changes: 56 additions & 0 deletions pandas/tests/indexes/datetimes/test_arithmetic.py
Original file line number Diff line number Diff line change
Expand Up @@ -571,6 +571,62 @@ def test_add_datetimelike_and_dti_tz(self, addend):
with tm.assert_raises_regex(TypeError, msg):
addend + dti_tz

# -------------------------------------------------------------
# __add__/__sub__ with ndarray[datetime64] and ndarray[timedelta64]

def test_dti_add_dt64_array_raises(self, tz):
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

these tests don't exist already? odd

dti = pd.date_range('2016-01-01', periods=3, tz=tz)
dtarr = dti.values

with pytest.raises(TypeError):
dti + dtarr
with pytest.raises(TypeError):
dtarr + dti

def test_dti_sub_dt64_array_naive(self):
dti = pd.date_range('2016-01-01', periods=3, tz=None)
dtarr = dti.values

expected = dti - dti
result = dti - dtarr
tm.assert_index_equal(result, expected)
result = dtarr - dti
tm.assert_index_equal(result, expected)

def test_dti_sub_dt64_array_aware_raises(self, tz):
if tz is None:
return
dti = pd.date_range('2016-01-01', periods=3, tz=tz)
dtarr = dti.values

with pytest.raises(TypeError):
dti - dtarr
with pytest.raises(TypeError):
dtarr - dti

def test_dti_add_td64_array(self, tz):
dti = pd.date_range('2016-01-01', periods=3, tz=tz)
tdi = dti - dti.shift(1)
tdarr = tdi.values

expected = dti + tdi
result = dti + tdarr
tm.assert_index_equal(result, expected)
result = tdarr + dti
tm.assert_index_equal(result, expected)

def test_dti_sub_td64_array(self, tz):
dti = pd.date_range('2016-01-01', periods=3, tz=tz)
tdi = dti - dti.shift(1)
tdarr = tdi.values

expected = dti - tdi
result = dti - tdarr
tm.assert_index_equal(result, expected)

with pytest.raises(TypeError):
tdarr - dti

# -------------------------------------------------------------

def test_sub_dti_dti(self):
Expand Down
58 changes: 58 additions & 0 deletions pandas/tests/indexes/period/test_arithmetic.py
Original file line number Diff line number Diff line change
Expand Up @@ -255,6 +255,64 @@ def test_comp_nat(self, dtype):


class TestPeriodIndexArithmetic(object):

# -----------------------------------------------------------------
# __add__/__sub__ with ndarray[datetime64] and ndarray[timedelta64]

def test_pi_add_sub_dt64_array_raises(self):
rng = pd.period_range('1/1/2000', freq='D', periods=3)
dti = pd.date_range('2016-01-01', periods=3)
dtarr = dti.values

with pytest.raises(TypeError):
rng + dtarr
with pytest.raises(TypeError):
dtarr + rng

with pytest.raises(TypeError):
rng - dtarr
with pytest.raises(TypeError):
dtarr - rng

def test_pi_add_sub_td64_array_non_tick_raises(self):
rng = pd.period_range('1/1/2000', freq='Q', periods=3)
dti = pd.date_range('2016-01-01', periods=3)
tdi = dti - dti.shift(1)
tdarr = tdi.values

with pytest.raises(period.IncompatibleFrequency):
rng + tdarr
with pytest.raises(period.IncompatibleFrequency):
tdarr + rng

with pytest.raises(period.IncompatibleFrequency):
rng - tdarr
with pytest.raises(period.IncompatibleFrequency):
tdarr - rng

@pytest.mark.xfail(reason='op with TimedeltaIndex raises, with ndarray OK')
def test_pi_add_sub_td64_array_tick(self):
rng = pd.period_range('1/1/2000', freq='Q', periods=3)
dti = pd.date_range('2016-01-01', periods=3)
tdi = dti - dti.shift(1)
tdarr = tdi.values

expected = rng + tdi
result = rng + tdarr
tm.assert_index_equal(result, expected)
result = tdarr + rng
tm.assert_index_equal(result, expected)

expected = rng - tdi
result = rng - tdarr
tm.assert_index_equal(result, expected)

with pytest.raises(TypeError):
tdarr - rng

# -----------------------------------------------------------------
# operations with array/Index of DateOffset objects

@pytest.mark.parametrize('box', [np.array, pd.Index])
def test_pi_add_offset_array(self, box):
# GH#18849
Expand Down
49 changes: 49 additions & 0 deletions pandas/tests/indexes/timedeltas/test_arithmetic.py
Original file line number Diff line number Diff line change
Expand Up @@ -494,6 +494,55 @@ def test_tdi_radd_timestamp(self):
expected = DatetimeIndex(['2011-01-02', '2011-01-03'])
tm.assert_index_equal(result, expected)

# -------------------------------------------------------------
# __add__/__sub__ with ndarray[datetime64] and ndarray[timedelta64]

def test_tdi_sub_dt64_array(self):
dti = pd.date_range('2016-01-01', periods=3)
tdi = dti - dti.shift(1)
dtarr = dti.values

with pytest.raises(TypeError):
tdi - dtarr

# TimedeltaIndex.__rsub__
expected = pd.DatetimeIndex(dtarr) - tdi
result = dtarr - tdi
tm.assert_index_equal(result, expected)

def test_tdi_add_dt64_array(self):
dti = pd.date_range('2016-01-01', periods=3)
tdi = dti - dti.shift(1)
dtarr = dti.values

expected = pd.DatetimeIndex(dtarr) + tdi
result = tdi + dtarr
tm.assert_index_equal(result, expected)
result = dtarr + tdi
tm.assert_index_equal(result, expected)

def test_tdi_add_td64_array(self):
dti = pd.date_range('2016-01-01', periods=3)
tdi = dti - dti.shift(1)
tdarr = tdi.values

expected = 2 * tdi
result = tdi + tdarr
tm.assert_index_equal(result, expected)
result = tdarr + tdi
tm.assert_index_equal(result, expected)

def test_tdi_sub_td64_array(self):
dti = pd.date_range('2016-01-01', periods=3)
tdi = dti - dti.shift(1)
tdarr = tdi.values

expected = 0 * tdi
result = tdi - tdarr
tm.assert_index_equal(result, expected)
result = tdarr - tdi
tm.assert_index_equal(result, expected)

# -------------------------------------------------------------

@pytest.mark.parametrize('scalar_td', [
Expand Down