Skip to content

Commit

Permalink
BUG: fix+test DTA/TDA/PA add/sub Index (#27726)
Browse files Browse the repository at this point in the history
  • Loading branch information
jbrockmendel authored and jreback committed Aug 5, 2019
1 parent 9fe8a0f commit 9b1c005
Show file tree
Hide file tree
Showing 4 changed files with 52 additions and 3 deletions.
6 changes: 3 additions & 3 deletions pandas/core/arrays/datetimelike.py
Original file line number Diff line number Diff line change
Expand Up @@ -1207,7 +1207,7 @@ def _time_shift(self, periods, freq=None):

def __add__(self, other):
other = lib.item_from_zerodim(other)
if isinstance(other, (ABCSeries, ABCDataFrame)):
if isinstance(other, (ABCSeries, ABCDataFrame, ABCIndexClass)):
return NotImplemented

# scalar others
Expand Down Expand Up @@ -1273,7 +1273,7 @@ def __radd__(self, other):

def __sub__(self, other):
other = lib.item_from_zerodim(other)
if isinstance(other, (ABCSeries, ABCDataFrame)):
if isinstance(other, (ABCSeries, ABCDataFrame, ABCIndexClass)):
return NotImplemented

# scalar others
Expand Down Expand Up @@ -1340,7 +1340,7 @@ def __sub__(self, other):
return result

def __rsub__(self, other):
if is_datetime64_dtype(other) and is_timedelta64_dtype(self):
if is_datetime64_any_dtype(other) and is_timedelta64_dtype(self):
# ndarray[datetime64] cannot be subtracted from self, so
# we need to wrap in DatetimeArray/Index and flip the operation
if not isinstance(other, DatetimeLikeArrayMixin):
Expand Down
18 changes: 18 additions & 0 deletions pandas/tests/arithmetic/test_datetime64.py
Original file line number Diff line number Diff line change
Expand Up @@ -2249,6 +2249,23 @@ def test_add_datetimelike_and_dti(self, addend, tz):

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

def test_dta_add_sub_index(self, tz_naive_fixture):
# Check that DatetimeArray defers to Index classes
dti = date_range("20130101", periods=3, tz=tz_naive_fixture)
dta = dti.array
result = dta - dti
expected = dti - dti
tm.assert_index_equal(result, expected)

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

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

def test_sub_dti_dti(self):
# previously performed setop (deprecated in 0.16.0), now changed to
# return subtraction -> TimeDeltaIndex (GH ...)
Expand Down Expand Up @@ -2554,6 +2571,7 @@ def test_shift_months(years, months):
tm.assert_index_equal(actual, expected)


# FIXME: this belongs in scalar tests
class SubDatetime(datetime):
pass

Expand Down
12 changes: 12 additions & 0 deletions pandas/tests/arithmetic/test_period.py
Original file line number Diff line number Diff line change
Expand Up @@ -1041,6 +1041,18 @@ def test_parr_add_sub_tdt64_nat_array(self, box_df_fail, other):
with pytest.raises(TypeError):
other - obj

# ---------------------------------------------------------------
# Unsorted

def test_parr_add_sub_index(self):
# Check that PeriodArray defers to Index on arithmetic ops
pi = pd.period_range("2000-12-31", periods=3)
parr = pi.array

result = parr - pi
expected = pi - pi
tm.assert_index_equal(result, expected)


class TestPeriodSeriesArithmetic:
def test_ops_series_timedelta(self):
Expand Down
19 changes: 19 additions & 0 deletions pandas/tests/arithmetic/test_timedelta64.py
Original file line number Diff line number Diff line change
Expand Up @@ -480,6 +480,25 @@ def test_timedelta(self, freq):
tm.assert_index_equal(result1, result4)
tm.assert_index_equal(result2, result3)

def test_tda_add_sub_index(self):
# Check that TimedeltaArray defers to Index on arithmetic ops
tdi = TimedeltaIndex(["1 days", pd.NaT, "2 days"])
tda = tdi.array

dti = pd.date_range("1999-12-31", periods=3, freq="D")

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

result = tda + tdi
expected = tdi + tdi
tm.assert_index_equal(result, expected)

result = tda - tdi
expected = tdi - tdi
tm.assert_index_equal(result, expected)


class TestAddSubNaTMasking:
# TODO: parametrize over boxes
Expand Down

0 comments on commit 9b1c005

Please sign in to comment.