REF: De-duplicate pieces of datetimelike arithmetic

pandas/core/arrays/datetimelike.py

@@ -101,6 +101,15 @@ def _shallow_copy(self, values=None, **kwargs):
             attributes['dtype'] = self.dtype
         return self._simple_new(values, **attributes)
 
+    def _semi_shallow_copy(self, values):
+        # similar to shallow_copy, but with freq="infer"
+        if is_period_dtype(self):
+            return self._shallow_copy(values)
+        elif is_timedelta64_dtype(self):
+            return type(self)(values, freq='infer')
+        else:
+            return type(self)(values, tz=self.tz, freq="infer")
+
 
 class DatetimeLikeArrayMixin(ExtensionOpsMixin, AttributesMixin):
     """
@@ -246,27 +255,6 @@ def _maybe_mask_results(self, result, fill_value=None, convert=None):
             result[self._isnan] = fill_value
         return result
 
-    def _nat_new(self, box=True):
-        """
-        Return Array/Index or ndarray filled with NaT which has the same
-        length as the caller.
-
-        Parameters
-        ----------
-        box : boolean, default True
-            - If True returns a Array/Index as the same as caller.
-            - If False returns ndarray of np.int64.
-        """
-        result = np.zeros(len(self), dtype=np.int64)
-        result.fill(iNaT)
-        if not box:
-            return result
-
-        attribs = self._get_attributes_dict()
-        if not is_period_dtype(self):
-            attribs['freq'] = None
-        return self._simple_new(result, **attribs)
-
     # ------------------------------------------------------------------
     # Frequency Properties/Methods
 
@@ -360,8 +348,38 @@ def _sub_period(self, other):
     def _add_offset(self, offset):
         raise com.AbstractMethodError(self)
 
-    def _add_delta(self, other):
-        return NotImplemented
+    def _add_delta(self, delta):
+        """
+        Add a timedelta-like, DateOffset, or TimedeltaIndex-like object
+        to self.
+
+        Parameters
+        ----------
+        delta : {timedelta, np.timedelta64, DateOffset,
+                 TimedeltaIndex, ndarray[timedelta64]}
+
+        Returns
+        -------
+        result : same type as self
+
+        Notes
+        -----
+        The result's name is set outside of _add_delta by the calling
+        method (__add__ or __sub__)
+        """
+        if is_period_dtype(self) and not isinstance(self.freq, Tick):
+            # We cannot add timedelta-like to non-tick PeriodArray
+            raise IncompatibleFrequency("Input has different freq from "
+                                        "{cls}(freq={freqstr})"
+                                        .format(cls=type(self).__name__,
+                                                freqstr=self.freqstr))
+
+        if isinstance(delta, (Tick, timedelta, np.timedelta64)):
+            new_values = self._add_delta_td(delta)
+        elif is_timedelta64_dtype(delta):
+            new_values = self._add_delta_tdi(delta)
+
+        return self._semi_shallow_copy(new_values)
 
     def _add_delta_td(self, other):
         """
@@ -371,16 +389,15 @@ def _add_delta_td(self, other):
         inc = delta_to_nanoseconds(other)
         new_values = checked_add_with_arr(self.asi8, inc,
                                           arr_mask=self._isnan).view('i8')
-        if self.hasnans:
-            new_values[self._isnan] = iNaT
+        new_values = self._maybe_mask_results(new_values, fill_value=iNaT)
         return new_values.view('i8')
 
     def _add_delta_tdi(self, other):
         """
         Add a delta of a TimedeltaIndex
         return the i8 result view
         """
-        if not len(self) == len(other):
+        if len(self) != len(other):
             raise ValueError("cannot add indices of unequal length")
 
         if isinstance(other, np.ndarray):
@@ -407,7 +424,9 @@ def _add_nat(self):
 
         # GH#19124 pd.NaT is treated like a timedelta for both timedelta
         # and datetime dtypes
-        return self._nat_new(box=True)
+        result = np.zeros(len(self), dtype=np.int64)
+        result.fill(iNaT)
+        return self._shallow_copy(result, freq=None)
 
     def _sub_nat(self):
         """Subtract pd.NaT from self"""
@@ -441,7 +460,7 @@ def _sub_period_array(self, other):
                             .format(dtype=other.dtype,
                                     cls=type(self).__name__))
 
-        if not len(self) == len(other):
+        if len(self) != len(other):
             raise ValueError("cannot subtract arrays/indices of "
                              "unequal length")
         if self.freq != other.freq:
@@ -481,7 +500,7 @@ def _addsub_int_array(self, other, op):
             res_values = checked_add_with_arr(self.asi8, other,
                                               arr_mask=self._isnan)
             res_values = res_values.view('i8')
-            res_values[self._isnan] = iNaT
+            res_values = self._maybe_mask_results(res_values, fill_value=iNaT)
             return self._from_ordinals(res_values, freq=self.freq)
 
         elif self.freq is None:

pandas/core/arrays/datetimes.py

@@ -425,10 +425,20 @@ def _assert_tzawareness_compat(self, other):
     # Arithmetic Methods
 
     def _sub_datelike_dti(self, other):
-        """subtraction of two DatetimeIndexes"""
-        if not len(self) == len(other):
+        """subtraction of DatetimeArray/Index or ndarray[datetime64]"""
+        if len(self) != len(other):
             raise ValueError("cannot add indices of unequal length")
 
+        if not isinstance(other, DatetimeArrayMixin):
+            # i.e. np.ndarray; we assume it is datetime64-dtype
+            other = type(self)(other)
+
+        if not self._has_same_tz(other):
+            # require tz compat
+            raise TypeError("{cls} subtraction must have the same "
+                            "timezones or no timezones"
+                            .format(cls=type(self).__name__))
+
         self_i8 = self.asi8
         other_i8 = other.asi8
         new_values = checked_add_with_arr(self_i8, -other_i8,
@@ -459,72 +469,24 @@ def _add_offset(self, offset):
     def _sub_datelike(self, other):
         # subtract a datetime from myself, yielding a ndarray[timedelta64[ns]]
         if isinstance(other, (DatetimeArrayMixin, np.ndarray)):
-            if isinstance(other, np.ndarray):
-                # if other is an ndarray, we assume it is datetime64-dtype
-                other = type(self)(other)
-            if not self._has_same_tz(other):
-                # require tz compat
-                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 NaT
-            other = Timestamp(other)
-            if other is NaT:
-                return self - NaT
+            return self._sub_datelike_dti(other)
+
+        assert isinstance(other, (datetime, np.datetime64))
+        assert other is not NaT
+        other = Timestamp(other)
+        if other is NaT:
+            return self - NaT
+        elif not self._has_same_tz(other):
             # require tz compat
-            elif not self._has_same_tz(other):
-                raise TypeError("Timestamp subtraction must have the same "
-                                "timezones or no timezones")
-            else:
-                i8 = self.asi8
-                result = checked_add_with_arr(i8, -other.value,
-                                              arr_mask=self._isnan)
-                result = self._maybe_mask_results(result,
-                                                  fill_value=iNaT)
-        else:
-            raise TypeError("cannot subtract {cls} and {typ}"
-                            .format(cls=type(self).__name__,
-                                    typ=type(other).__name__))
+            raise TypeError("Timestamp subtraction must have the same "
+                            "timezones or no timezones")
+        i8 = self.asi8
+        result = checked_add_with_arr(i8, -other.value,
+                                      arr_mask=self._isnan)
+        result = self._maybe_mask_results(result,
+                                          fill_value=iNaT)
         return result.view('timedelta64[ns]')
 
-    def _add_delta(self, delta):
-        """
-        Add a timedelta-like, DateOffset, or TimedeltaIndex-like object
-        to self.
-
-        Parameters
-        ----------
-        delta : {timedelta, np.timedelta64, DateOffset,
-                 TimedeltaIndex, ndarray[timedelta64]}
-
-        Returns
-        -------
-        result : same type as self
-
-        Notes
-        -----
-        The result's name is set outside of _add_delta by the calling
-        method (__add__ or __sub__)
-        """
-        from pandas.core.arrays.timedeltas import TimedeltaArrayMixin
-
-        if isinstance(delta, (Tick, timedelta, np.timedelta64)):
-            new_values = self._add_delta_td(delta)
-        elif is_timedelta64_dtype(delta):
-            if not isinstance(delta, TimedeltaArrayMixin):
-                delta = TimedeltaArrayMixin(delta)
-            new_values = self._add_delta_tdi(delta)
-        else:
-            new_values = self.astype('O') + delta
-
-        tz = 'UTC' if self.tz is not None else None
-        result = type(self)(new_values, tz=tz, freq='infer')
-        if self.tz is not None and self.tz is not utc:
-            result = result.tz_convert(self.tz)
-        return result
-
     # -----------------------------------------------------------------
     # Timezone Conversion and Localization Methods
 

pandas/core/arrays/period.py

@@ -5,7 +5,6 @@
 
 import numpy as np
 
-from pandas._libs import lib
 from pandas._libs.tslib import NaT, iNaT
 from pandas._libs.tslibs.period import (
     Period, IncompatibleFrequency, DIFFERENT_FREQ_INDEX,
@@ -26,7 +25,7 @@
 import pandas.core.common as com
 
 from pandas.tseries import frequencies
-from pandas.tseries.offsets import Tick, DateOffset
+from pandas.tseries.offsets import Tick
 
 from pandas.core.arrays import datetimelike as dtl
 from pandas.core.arrays.datetimelike import DatetimeLikeArrayMixin
@@ -379,39 +378,6 @@ def _add_delta_tdi(self, other):
         delta = self._check_timedeltalike_freq_compat(other)
         return self._addsub_int_array(delta, operator.add)
 
-    def _add_delta(self, other):
-        """
-        Add a timedelta-like, Tick, or TimedeltaIndex-like object
-        to self.
-
-        Parameters
-        ----------
-        other : {timedelta, np.timedelta64, Tick,
-                 TimedeltaIndex, ndarray[timedelta64]}
-
-        Returns
-        -------
-        result : same type as self
-        """
-        if not isinstance(self.freq, Tick):
-            # We cannot add timedelta-like to non-tick PeriodArray
-            raise IncompatibleFrequency("Input has different freq from "
-                                        "{cls}(freq={freqstr})"
-                                        .format(cls=type(self).__name__,
-                                                freqstr=self.freqstr))
-
-        # TODO: standardize across datetimelike subclasses whether to return
-        #  i8 view or _shallow_copy
-        if isinstance(other, (Tick, timedelta, np.timedelta64)):
-            new_values = self._add_delta_td(other)
-            return self._shallow_copy(new_values)
-        elif is_timedelta64_dtype(other):
-            # ndarray[timedelta64] or TimedeltaArray/index
-            new_values = self._add_delta_tdi(other)
-            return self._shallow_copy(new_values)
-        else:  # pragma: no cover
-            raise TypeError(type(other).__name__)
-
     @deprecate_kwarg(old_arg_name='n', new_arg_name='periods')
     def shift(self, periods):
         """
@@ -445,48 +411,6 @@ def _time_shift(self, n):
             values[self._isnan] = iNaT
         return self._shallow_copy(values=values)
 
-    def _maybe_convert_timedelta(self, other):
-        """
-        Convert timedelta-like input to an integer multiple of self.freq
-
-        Parameters
-        ----------
-        other : timedelta, np.timedelta64, DateOffset, int, np.ndarray
-
-        Returns
-        -------
-        converted : int, np.ndarray[int64]
-
-        Raises
-        ------
-        IncompatibleFrequency : if the input cannot be written as a multiple
-            of self.freq.  Note IncompatibleFrequency subclasses ValueError.
-        """
-        if isinstance(
-                other, (timedelta, np.timedelta64, Tick, np.ndarray)):
-            offset = frequencies.to_offset(self.freq.rule_code)
-            if isinstance(offset, Tick):
-                # _check_timedeltalike_freq_compat will raise if incompatible
-                delta = self._check_timedeltalike_freq_compat(other)
-                return delta
-        elif isinstance(other, DateOffset):
-            freqstr = other.rule_code
-            base = frequencies.get_base_alias(freqstr)
-            if base == self.freq.rule_code:
-                return other.n
-            msg = DIFFERENT_FREQ_INDEX.format(self.freqstr, other.freqstr)
-            raise IncompatibleFrequency(msg)
-        elif lib.is_integer(other):
-            # integer is passed to .shift via
-            # _add_datetimelike_methods basically
-            # but ufunc may pass integer to _add_delta
-            return other
-
-        # raise when input doesn't have freq
-        msg = "Input has different freq from {cls}(freq={freqstr})"
-        raise IncompatibleFrequency(msg.format(cls=type(self).__name__,
-                                               freqstr=self.freqstr))
-
     def _check_timedeltalike_freq_compat(self, other):
         """
         Arithmetic operations with timedelta-like scalars or array `other`

pandas/core/arrays/timedeltas.py

@@ -196,38 +196,6 @@ def _sub_datelike(self, other):
         raise TypeError("cannot subtract a datelike from a {cls}"
                         .format(cls=type(self).__name__))
 
-    def _add_delta(self, delta):
-        """
-        Add a timedelta-like, Tick, or TimedeltaIndex-like object
-        to self.
-
-        Parameters
-        ----------
-        delta : timedelta, np.timedelta64, Tick, TimedeltaArray, TimedeltaIndex
-
-        Returns
-        -------
-        result : same type as self
-
-        Notes
-        -----
-        The result's name is set outside of _add_delta by the calling
-        method (__add__ or __sub__)
-        """
-        if isinstance(delta, (Tick, timedelta, np.timedelta64)):
-            new_values = self._add_delta_td(delta)
-        elif isinstance(delta, TimedeltaArrayMixin):
-            new_values = self._add_delta_tdi(delta)
-        elif is_timedelta64_dtype(delta):
-            # ndarray[timedelta64] --> wrap in TimedeltaArray/Index
-            delta = type(self)(delta)
-            new_values = self._add_delta_tdi(delta)
-        else:
-            raise TypeError("cannot add the type {0} to a TimedeltaIndex"
-                            .format(type(delta)))
-
-        return type(self)(new_values, freq='infer')
-
     def _add_datelike(self, other):
         # adding a timedeltaindex to a datetimelike
         from pandas.core.arrays import DatetimeArrayMixin