BUG: fix+test DTA/TDA/PA add/sub Index (pandas-dev#27726)

Author: jbrockmendel

pandas/core/arrays/datetimelike.py

@@ -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
@@ -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
@@ -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):

pandas/tests/arithmetic/test_datetime64.py

@@ -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 ...)
@@ -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
 

pandas/tests/arithmetic/test_period.py

@@ -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):

pandas/tests/arithmetic/test_timedelta64.py

@@ -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