BUG/TST/REF: Datetimelike Arithmetic Methods (#23215)

Author: jbrockmendel

doc/source/whatsnew/v0.24.0.txt

@@ -1028,6 +1028,7 @@ Datetimelike
 - Bug in :func:`date_range` when decrementing a start date to a past end date by a negative frequency (:issue:`23270`)
 - Bug in :func:`DataFrame.combine` with datetimelike values raising a TypeError (:issue:`23079`)
 - Bug in :func:`date_range` with frequency of ``Day`` or higher where dates sufficiently far in the future could wrap around to the past instead of raising ``OutOfBoundsDatetime`` (:issue:`14187`)
+- Bug in :class:`PeriodIndex` with attribute ``freq.n`` greater than 1 where adding a :class:`DateOffset` object would return incorrect results (:issue:`23215`)
 
 Timedelta
 ^^^^^^^^^
@@ -1039,7 +1040,8 @@ Timedelta
 - Bug in :class:`TimedeltaIndex` incorrectly allowing indexing with ``Timestamp`` object (:issue:`20464`)
 - Fixed bug where subtracting :class:`Timedelta` from an object-dtyped array would raise ``TypeError`` (:issue:`21980`)
 - Fixed bug in adding a :class:`DataFrame` with all-`timedelta64[ns]` dtypes to a :class:`DataFrame` with all-integer dtypes returning incorrect results instead of raising ``TypeError`` (:issue:`22696`)
-
+- Bug in :class:`TimedeltaIndex` where adding a timezone-aware datetime scalar incorrectly returned a timezone-naive :class:`DatetimeIndex` (:issue:`23215`)
+- Bug in :class:`TimedeltaIndex` where adding ``np.timedelta64('NaT')`` incorrectly returned an all-`NaT` :class:`DatetimeIndex` instead of an all-`NaT` :class:`TimedeltaIndex` (:issue:`23215`)
 
 Timezones
 ^^^^^^^^^
@@ -1069,7 +1071,7 @@ Offsets
 
 - Bug in :class:`FY5253` where date offsets could incorrectly raise an ``AssertionError`` in arithmetic operatons (:issue:`14774`)
 - Bug in :class:`DateOffset` where keyword arguments ``week`` and ``milliseconds`` were accepted and ignored.  Passing these will now raise ``ValueError`` (:issue:`19398`)
--
+- Bug in adding :class:`DateOffset` with :class:`DataFrame` or :class:`PeriodIndex` incorrectly raising ``TypeError`` (:issue:`23215`)
 
 Numeric
 ^^^^^^^

pandas/_libs/tslibs/offsets.pyx

@@ -344,8 +344,8 @@ class _BaseOffset(object):
         return {name: kwds[name] for name in kwds if kwds[name] is not None}
 
     def __add__(self, other):
-        if getattr(other, "_typ", None) in ["datetimeindex",
-                                            "series", "period"]:
+        if getattr(other, "_typ", None) in ["datetimeindex", "periodindex",
+                                            "series", "period", "dataframe"]:
             # defer to the other class's implementation
             return other + self
         try:

pandas/core/arrays/datetimelike.py

@@ -221,11 +221,12 @@ def hasnans(self):
         """ return if I have any nans; enables various perf speedups """
         return bool(self._isnan.any())
 
-    def _maybe_mask_results(self, result, fill_value=None, convert=None):
+    def _maybe_mask_results(self, result, fill_value=iNaT, convert=None):
         """
         Parameters
         ----------
         result : a ndarray
+        fill_value : object, default iNaT
         convert : string/dtype or None
 
         Returns
@@ -246,27 +247,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
 
@@ -346,41 +326,74 @@ def _validate_frequency(cls, index, freq, **kwargs):
     # ------------------------------------------------------------------
     # Arithmetic Methods
 
-    def _add_datelike(self, other):
+    def _add_datetimelike_scalar(self, other):
+        # Overriden by TimedeltaArray
         raise TypeError("cannot add {cls} and {typ}"
                         .format(cls=type(self).__name__,
                                 typ=type(other).__name__))
 
-    def _sub_datelike(self, other):
-        raise com.AbstractMethodError(self)
+    _add_datetime_arraylike = _add_datetimelike_scalar
+
+    def _sub_datetimelike_scalar(self, other):
+        # Overridden by DatetimeArray
+        assert other is not NaT
+        raise TypeError("cannot subtract a datelike from a {cls}"
+                        .format(cls=type(self).__name__))
+
+    _sub_datetime_arraylike = _sub_datetimelike_scalar
 
     def _sub_period(self, other):
-        return NotImplemented
+        # Overriden by PeriodArray
+        raise TypeError("cannot subtract Period from a {cls}"
+                        .format(cls=type(self).__name__))
 
     def _add_offset(self, offset):
         raise com.AbstractMethodError(self)
 
     def _add_delta(self, other):
-        return NotImplemented
+        """
+        Add a timedelta-like, Tick or TimedeltaIndex-like object
+        to self, yielding an int64 numpy array
+
+        Parameters
+        ----------
+        delta : {timedelta, np.timedelta64, Tick,
+                 TimedeltaIndex, ndarray[timedelta64]}
+
+        Returns
+        -------
+        result : ndarray[int64]
 
-    def _add_delta_td(self, other):
+        Notes
+        -----
+        The result's name is set outside of _add_delta by the calling
+        method (__add__ or __sub__), if necessary (i.e. for Indexes).
+        """
+        if isinstance(other, (Tick, timedelta, np.timedelta64)):
+            new_values = self._add_timedeltalike_scalar(other)
+        elif is_timedelta64_dtype(other):
+            # ndarray[timedelta64] or TimedeltaArray/index
+            new_values = self._add_delta_tdi(other)
+
+        return new_values
+
+    def _add_timedeltalike_scalar(self, other):
         """
         Add a delta of a timedeltalike
         return the i8 result view
         """
         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)
         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 +420,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 +456,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:
@@ -473,6 +488,8 @@ def _addsub_int_array(self, other, op):
         -------
         result : same class as self
         """
+        # _addsub_int_array is overriden by PeriodArray
+        assert not is_period_dtype(self)
         assert op in [operator.add, operator.sub]
 
         if self.freq is None:
@@ -613,7 +630,7 @@ def __add__(self, other):
                 # specifically _not_ a Tick
                 result = self._add_offset(other)
             elif isinstance(other, (datetime, np.datetime64)):
-                result = self._add_datelike(other)
+                result = self._add_datetimelike_scalar(other)
             elif lib.is_integer(other):
                 # This check must come after the check for np.timedelta64
                 # as is_integer returns True for these
@@ -628,7 +645,7 @@ def __add__(self, other):
                 result = self._addsub_offset_array(other, operator.add)
             elif is_datetime64_dtype(other) or is_datetime64tz_dtype(other):
                 # DatetimeIndex, ndarray[datetime64]
-                return self._add_datelike(other)
+                return self._add_datetime_arraylike(other)
             elif is_integer_dtype(other):
                 result = self._addsub_int_array(other, operator.add)
             elif is_float_dtype(other):
@@ -671,7 +688,7 @@ def __sub__(self, other):
                 # specifically _not_ a Tick
                 result = self._add_offset(-other)
             elif isinstance(other, (datetime, np.datetime64)):
-                result = self._sub_datelike(other)
+                result = self._sub_datetimelike_scalar(other)
             elif lib.is_integer(other):
                 # This check must come after the check for np.timedelta64
                 # as is_integer returns True for these
@@ -688,7 +705,7 @@ def __sub__(self, other):
                 result = self._addsub_offset_array(other, operator.sub)
             elif is_datetime64_dtype(other) or is_datetime64tz_dtype(other):
                 # DatetimeIndex, ndarray[datetime64]
-                result = self._sub_datelike(other)
+                result = self._sub_datetime_arraylike(other)
             elif is_period_dtype(other):
                 # PeriodIndex
                 result = self._sub_period_array(other)

pandas/core/arrays/datetimes.py

@@ -1,5 +1,5 @@
 # -*- coding: utf-8 -*-
-from datetime import datetime, timedelta, time
+from datetime import datetime, time
 import warnings
 
 import numpy as np
@@ -21,7 +21,6 @@
     is_object_dtype,
     is_datetime64tz_dtype,
     is_datetime64_dtype,
-    is_timedelta64_dtype,
     ensure_int64)
 from pandas.core.dtypes.dtypes import DatetimeTZDtype
 from pandas.core.dtypes.missing import isna
@@ -76,11 +75,12 @@ def f(self):
 
         if field in self._object_ops:
             result = fields.get_date_name_field(values, field)
-            result = self._maybe_mask_results(result)
+            result = self._maybe_mask_results(result, fill_value=None)
 
         else:
             result = fields.get_date_field(values, field)
-            result = self._maybe_mask_results(result, convert='float64')
+            result = self._maybe_mask_results(result, fill_value=None,
+                                              convert='float64')
 
         return result
 
@@ -424,11 +424,21 @@ def _assert_tzawareness_compat(self, other):
     # -----------------------------------------------------------------
     # Arithmetic Methods
 
-    def _sub_datelike_dti(self, other):
-        """subtraction of two DatetimeIndexes"""
-        if not len(self) == len(other):
+    def _sub_datetime_arraylike(self, other):
+        """subtract DatetimeArray/Index or ndarray[datetime64]"""
+        if len(self) != len(other):
             raise ValueError("cannot add indices of unequal length")
 
+        if isinstance(other, np.ndarray):
+            assert is_datetime64_dtype(other)
+            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,
@@ -456,74 +466,41 @@ def _add_offset(self, offset):
 
         return type(self)(result, freq='infer')
 
-    def _sub_datelike(self, other):
+    def _sub_datetimelike_scalar(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
+        assert isinstance(other, (datetime, np.datetime64))
+        assert other is not NaT
+        other = Timestamp(other)
+        if other is NaT:
+            return self - NaT
+
+        if 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)
         return result.view('timedelta64[ns]')
 
     def _add_delta(self, delta):
         """
-        Add a timedelta-like, DateOffset, or TimedeltaIndex-like object
-        to self.
+        Add a timedelta-like, Tick, or TimedeltaIndex-like object
+        to self, yielding a new DatetimeArray
 
         Parameters
         ----------
-        delta : {timedelta, np.timedelta64, DateOffset,
+        other : {timedelta, np.timedelta64, Tick,
                  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__)
+        result : DatetimeArray
         """
-        from pandas.core.arrays 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
+        new_values = dtl.DatetimeLikeArrayMixin._add_delta(self, delta)
+        return type(self)(new_values, tz=self.tz, freq='infer')
 
     # -----------------------------------------------------------------
     # Timezone Conversion and Localization Methods
@@ -904,7 +881,7 @@ def month_name(self, locale=None):
 
         result = fields.get_date_name_field(values, 'month_name',
                                             locale=locale)
-        result = self._maybe_mask_results(result)
+        result = self._maybe_mask_results(result, fill_value=None)
         return result
 
     def day_name(self, locale=None):
@@ -940,7 +917,7 @@ def day_name(self, locale=None):
 
         result = fields.get_date_name_field(values, 'day_name',
                                             locale=locale)
-        result = self._maybe_mask_results(result)
+        result = self._maybe_mask_results(result, fill_value=None)
         return result
 
     @property