Fix+test timedelta64(nat) ops (#23425)

Follow-up to gh-23320.

Author: jbrockmendel

pandas/core/arrays/datetimelike.py

@@ -36,6 +36,7 @@
     is_object_dtype)
 from pandas.core.dtypes.generic import ABCSeries, ABCDataFrame, ABCIndexClass
 from pandas.core.dtypes.dtypes import DatetimeTZDtype
+from pandas.core.dtypes.missing import isna
 
 import pandas.core.common as com
 from pandas.core.algorithms import checked_add_with_arr
@@ -370,6 +371,12 @@ def _add_timedeltalike_scalar(self, other):
         Add a delta of a timedeltalike
         return the i8 result view
         """
+        if isna(other):
+            # i.e np.timedelta64("NaT"), not recognized by delta_to_nanoseconds
+            new_values = np.empty(len(self), dtype='i8')
+            new_values[:] = iNaT
+            return new_values
+
         inc = delta_to_nanoseconds(other)
         new_values = checked_add_with_arr(self.asi8, inc,
                                           arr_mask=self._isnan).view('i8')
@@ -442,7 +449,7 @@ def _sub_period_array(self, other):
             Array of DateOffset objects; nulls represented by NaT
         """
         if not is_period_dtype(self):
-            raise TypeError("cannot subtract {dtype}-dtype to {cls}"
+            raise TypeError("cannot subtract {dtype}-dtype from {cls}"
                             .format(dtype=other.dtype,
                                     cls=type(self).__name__))
 
@@ -741,6 +748,11 @@ def __rsub__(self, other):
                 raise TypeError("cannot subtract {cls} from {typ}"
                                 .format(cls=type(self).__name__,
                                         typ=type(other).__name__))
+            elif is_period_dtype(self) and is_timedelta64_dtype(other):
+                # TODO: Can we simplify/generalize these cases at all?
+                raise TypeError("cannot subtract {cls} from {dtype}"
+                                .format(cls=type(self).__name__,
+                                        dtype=other.dtype))
             return -(self - other)
         cls.__rsub__ = __rsub__
 

pandas/core/arrays/period.py

@@ -35,7 +35,7 @@
 from pandas.core.dtypes.generic import (
     ABCSeries, ABCIndexClass, ABCPeriodIndex
 )
-from pandas.core.dtypes.missing import isna
+from pandas.core.dtypes.missing import isna, notna
 from pandas.core.missing import pad_1d, backfill_1d
 
 import pandas.core.common as com
@@ -149,6 +149,8 @@ class PeriodArray(dtl.DatetimeLikeArrayMixin, ExtensionArray):
     period_array : Create a new PeriodArray
     pandas.PeriodIndex : Immutable Index for period data
     """
+    # array priority higher than numpy scalars
+    __array_priority__ = 1000
     _attributes = ["freq"]
     _typ = "periodarray"  # ABCPeriodArray
 
@@ -761,12 +763,15 @@ def _add_timedeltalike_scalar(self, other):
         assert isinstance(self.freq, Tick)  # checked by calling function
         assert isinstance(other, (timedelta, np.timedelta64, Tick))
 
-        delta = self._check_timedeltalike_freq_compat(other)
+        if notna(other):
+            # special handling for np.timedelta64("NaT"), avoid calling
+            #  _check_timedeltalike_freq_compat as that would raise TypeError
+            other = self._check_timedeltalike_freq_compat(other)
 
         # Note: when calling parent class's _add_timedeltalike_scalar,
         #  it will call delta_to_nanoseconds(delta).  Because delta here
         #  is an integer, delta_to_nanoseconds will return it unchanged.
-        ordinals = super(PeriodArray, self)._add_timedeltalike_scalar(delta)
+        ordinals = super(PeriodArray, self)._add_timedeltalike_scalar(other)
         return ordinals
 
     def _add_delta_tdi(self, other):

pandas/core/indexes/base.py

@@ -4707,7 +4707,7 @@ def _evaluate_with_timedelta_like(self, other, op):
                                                       'radd', 'rsub']:
             raise TypeError("Operation {opname} between {cls} and {other} "
                             "is invalid".format(opname=op.__name__,
-                                                cls=type(self).__name__,
+                                                cls=self.dtype,
                                                 other=type(other).__name__))
 
         other = Timedelta(other)

pandas/tests/arithmetic/test_datetime64.py

@@ -1188,6 +1188,25 @@ def test_dti_isub_timedeltalike(self, tz_naive_fixture, two_hours):
         rng -= two_hours
         tm.assert_index_equal(rng, expected)
 
+    def test_dt64arr_add_sub_td64_nat(self, box, tz_naive_fixture):
+        # GH#23320 special handling for timedelta64("NaT")
+        tz = tz_naive_fixture
+        dti = pd.date_range("1994-04-01", periods=9, tz=tz, freq="QS")
+        other = np.timedelta64("NaT")
+        expected = pd.DatetimeIndex(["NaT"] * 9, tz=tz)
+
+        obj = tm.box_expected(dti, box)
+        expected = tm.box_expected(expected, box)
+
+        result = obj + other
+        tm.assert_equal(result, expected)
+        result = other + obj
+        tm.assert_equal(result, expected)
+        result = obj - other
+        tm.assert_equal(result, expected)
+        with pytest.raises(TypeError):
+            other - obj
+
     # -------------------------------------------------------------
     # Binary operations DatetimeIndex and TimedeltaIndex/array
     def test_dti_add_tdi(self, tz_naive_fixture):

pandas/tests/arithmetic/test_period.py

@@ -419,7 +419,7 @@ def test_pi_add_sub_td64_array_non_tick_raises(self):
 
         with pytest.raises(period.IncompatibleFrequency):
             rng - tdarr
-        with pytest.raises(period.IncompatibleFrequency):
+        with pytest.raises(TypeError):
             tdarr - rng
 
     def test_pi_add_sub_td64_array_tick(self):
@@ -801,6 +801,24 @@ def test_pi_add_sub_timedeltalike_freq_mismatch_monthly(self,
         with tm.assert_raises_regex(period.IncompatibleFrequency, msg):
             rng -= other
 
+    def test_parr_add_sub_td64_nat(self, box):
+        # GH#23320 special handling for timedelta64("NaT")
+        pi = pd.period_range("1994-04-01", periods=9, freq="19D")
+        other = np.timedelta64("NaT")
+        expected = pd.PeriodIndex(["NaT"] * 9, freq="19D")
+
+        obj = tm.box_expected(pi, box)
+        expected = tm.box_expected(expected, box)
+
+        result = obj + other
+        tm.assert_equal(result, expected)
+        result = other + obj
+        tm.assert_equal(result, expected)
+        result = obj - other
+        tm.assert_equal(result, expected)
+        with pytest.raises(TypeError):
+            other - obj
+
 
 class TestPeriodSeriesArithmetic(object):
     def test_ops_series_timedelta(self):

pandas/tests/arithmetic/test_timedelta64.py

@@ -735,6 +735,24 @@ def test_td64arr_add_sub_tdi(self, box_df_broadcast_failure, names):
         else:
             assert result.dtypes[0] == 'timedelta64[ns]'
 
+    def test_td64arr_add_sub_td64_nat(self, box):
+        # GH#23320 special handling for timedelta64("NaT")
+        tdi = pd.TimedeltaIndex([NaT, Timedelta('1s')])
+        other = np.timedelta64("NaT")
+        expected = pd.TimedeltaIndex(["NaT"] * 2)
+
+        obj = tm.box_expected(tdi, box)
+        expected = tm.box_expected(expected, box)
+
+        result = obj + other
+        tm.assert_equal(result, expected)
+        result = other + obj
+        tm.assert_equal(result, expected)
+        result = obj - other
+        tm.assert_equal(result, expected)
+        result = other - obj
+        tm.assert_equal(result, expected)
+
     def test_td64arr_sub_NaT(self, box):
         # GH#18808
         ser = Series([NaT, Timedelta('1s')])