handle NaT add/sub in one place (#19903)

Author: jbrockmendel

pandas/core/indexes/datetimelike.py

@@ -1,3 +1,4 @@
+# -*- coding: utf-8 -*-
 """
 Base and utility classes for tseries type pandas objects.
 """
@@ -640,6 +641,28 @@ def _add_datelike(self, other):
     def _sub_datelike(self, other):
         raise com.AbstractMethodError(self)
 
+    def _add_nat(self):
+        """Add pd.NaT to self"""
+        if is_period_dtype(self):
+            raise TypeError('Cannot add {cls} and {typ}'
+                            .format(cls=type(self).__name__,
+                                    typ=type(NaT).__name__))
+
+        # GH#19124 pd.NaT is treated like a timedelta for both timedelta
+        # and datetime dtypes
+        return self._nat_new(box=True)
+
+    def _sub_nat(self):
+        """Subtract pd.NaT from self"""
+        # GH#19124 Timedelta - datetime is not in general well-defined.
+        # We make an exception for pd.NaT, which in this case quacks
+        # like a timedelta.
+        # For datetime64 dtypes by convention we treat NaT as a datetime, so
+        # this subtraction returns a timedelta64 dtype.
+        # For period dtype, timedelta64 is a close-enough return dtype.
+        result = self._nat_new(box=False)
+        return result.view('timedelta64[ns]')
+
     def _sub_period(self, other):
         return NotImplemented
 
@@ -686,6 +709,8 @@ def __add__(self, other):
                 return NotImplemented
 
             # scalar others
+            elif other is NaT:
+                result = self._add_nat()
             elif isinstance(other, (DateOffset, timedelta, np.timedelta64)):
                 result = self._add_delta(other)
             elif isinstance(other, (datetime, np.datetime64)):
@@ -711,9 +736,13 @@ def __add__(self, other):
             else:  # pragma: no cover
                 return NotImplemented
 
-            if result is not NotImplemented:
-                res_name = ops.get_op_result_name(self, other)
-                result.name = res_name
+            if result is NotImplemented:
+                return NotImplemented
+            elif not isinstance(result, Index):
+                # Index.__new__ will choose appropriate subclass for dtype
+                result = Index(result)
+            res_name = ops.get_op_result_name(self, other)
+            result.name = res_name
             return result
 
         cls.__add__ = __add__
@@ -731,6 +760,8 @@ def __sub__(self, other):
                 return NotImplemented
 
             # scalar others
+            elif other is NaT:
+                result = self._sub_nat()
             elif isinstance(other, (DateOffset, timedelta, np.timedelta64)):
                 result = self._add_delta(-other)
             elif isinstance(other, (datetime, np.datetime64)):
@@ -762,9 +793,13 @@ def __sub__(self, other):
             else:  # pragma: no cover
                 return NotImplemented
 
-            if result is not NotImplemented:
-                res_name = ops.get_op_result_name(self, other)
-                result.name = res_name
+            if result is NotImplemented:
+                return NotImplemented
+            elif not isinstance(result, Index):
+                # Index.__new__ will choose appropriate subclass for dtype
+                result = Index(result)
+            res_name = ops.get_op_result_name(self, other)
+            result.name = res_name
             return result
 
         cls.__sub__ = __sub__

pandas/core/indexes/datetimes.py

@@ -853,32 +853,22 @@ def __setstate__(self, state):
             raise Exception("invalid pickle state")
     _unpickle_compat = __setstate__
 
-    def _add_datelike(self, other):
-        # adding a timedeltaindex to a datetimelike
-        if other is libts.NaT:
-            return self._nat_new(box=True)
-        raise TypeError("cannot add {0} and {1}"
-                        .format(type(self).__name__,
-                                type(other).__name__))
-
     def _sub_datelike(self, other):
-        # subtract a datetime from myself, yielding a TimedeltaIndex
-        from pandas import TimedeltaIndex
-
+        # subtract a datetime from myself, yielding a ndarray[timedelta64[ns]]
         if isinstance(other, (DatetimeIndex, np.ndarray)):
             # if other is an ndarray, we assume it is datetime64-dtype
             other = DatetimeIndex(other)
-
             # require tz compat
             if not self._has_same_tz(other):
                 raise TypeError("{cls} subtraction must have the same "
                                 "timezones or no timezones"
                                 .format(cls=type(self).__name__))
             result = self._sub_datelike_dti(other)
         elif isinstance(other, (datetime, np.datetime64)):
+            assert other is not libts.NaT
             other = Timestamp(other)
             if other is libts.NaT:
-                result = self._nat_new(box=False)
+                return self - libts.NaT
             # require tz compat
             elif not self._has_same_tz(other):
                 raise TypeError("Timestamp subtraction must have the same "
@@ -893,7 +883,7 @@ def _sub_datelike(self, other):
             raise TypeError("cannot subtract {cls} and {typ}"
                             .format(cls=type(self).__name__,
                                     typ=type(other).__name__))
-        return TimedeltaIndex(result)
+        return result.view('timedelta64[ns]')
 
     def _sub_datelike_dti(self, other):
         """subtraction of two DatetimeIndexes"""
@@ -906,7 +896,7 @@ def _sub_datelike_dti(self, other):
         if self.hasnans or other.hasnans:
             mask = (self._isnan) | (other._isnan)
             new_values[mask] = libts.iNaT
-        return new_values.view('i8')
+        return new_values.view('timedelta64[ns]')
 
     def _maybe_update_attributes(self, attrs):
         """ Update Index attributes (e.g. freq) depending on op """

pandas/core/indexes/period.py

@@ -12,7 +12,6 @@
     is_scalar,
     is_datetime64_dtype,
     is_datetime64_any_dtype,
-    is_timedelta64_dtype,
     is_period_dtype,
     is_bool_dtype,
     pandas_dtype,
@@ -23,7 +22,6 @@
 import pandas.tseries.frequencies as frequencies
 from pandas.tseries.frequencies import get_freq_code as _gfc
 from pandas.core.indexes.datetimes import DatetimeIndex, Int64Index, Index
-from pandas.core.indexes.timedeltas import TimedeltaIndex
 from pandas.core.indexes.datetimelike import DatelikeOps, DatetimeIndexOpsMixin
 from pandas.core.tools.datetimes import parse_time_string
 import pandas.tseries.offsets as offsets
@@ -700,16 +698,6 @@ def _maybe_convert_timedelta(self, other):
                 return other.n
             msg = _DIFFERENT_FREQ_INDEX.format(self.freqstr, other.freqstr)
             raise IncompatibleFrequency(msg)
-        elif isinstance(other, np.ndarray):
-            if is_integer_dtype(other):
-                return other
-            elif is_timedelta64_dtype(other):
-                offset = frequencies.to_offset(self.freq)
-                if isinstance(offset, offsets.Tick):
-                    nanos = delta_to_nanoseconds(other)
-                    offset_nanos = delta_to_nanoseconds(offset)
-                    if (nanos % offset_nanos).all() == 0:
-                        return nanos // offset_nanos
         elif is_integer(other):
             # integer is passed to .shift via
             # _add_datetimelike_methods basically
@@ -724,10 +712,7 @@ def _add_delta(self, other):
         return self.shift(ordinal_delta)
 
     def _sub_datelike(self, other):
-        if other is tslib.NaT:
-            new_data = np.empty(len(self), dtype=np.int64)
-            new_data.fill(tslib.iNaT)
-            return TimedeltaIndex(new_data)
+        assert other is not tslib.NaT
         return NotImplemented
 
     def _sub_period(self, other):

pandas/core/indexes/timedeltas.py

@@ -414,16 +414,13 @@ def _evaluate_with_timedelta_like(self, other, op):
     def _add_datelike(self, other):
         # adding a timedeltaindex to a datetimelike
         from pandas import Timestamp, DatetimeIndex
-
-        if other is NaT:
-            # GH#19124 pd.NaT is treated like a timedelta
-            return self._nat_new()
-        elif isinstance(other, (DatetimeIndex, np.ndarray)):
+        if isinstance(other, (DatetimeIndex, np.ndarray)):
             # if other is an ndarray, we assume it is datetime64-dtype
             # defer to implementation in DatetimeIndex
             other = DatetimeIndex(other)
             return other + self
         else:
+            assert other is not NaT
             other = Timestamp(other)
             i8 = self.asi8
             result = checked_add_with_arr(i8, other.value,
@@ -432,14 +429,9 @@ def _add_datelike(self, other):
             return DatetimeIndex(result)
 
     def _sub_datelike(self, other):
-        # GH#19124 Timedelta - datetime is not in general well-defined.
-        # We make an exception for pd.NaT, which in this case quacks
-        # like a timedelta.
-        if other is NaT:
-            return self._nat_new()
-        else:
-            raise TypeError("cannot subtract a datelike from a {cls}"
-                            .format(cls=type(self).__name__))
+        assert other is not NaT
+        raise TypeError("cannot subtract a datelike from a {cls}"
+                        .format(cls=type(self).__name__))
 
     def _addsub_offset_array(self, other, op):
         # Add or subtract Array-like of DateOffset objects