REF: share DTA/TDA/PA arithmetic methods

pandas/_libs/tslibs/period.pyx

@@ -1728,10 +1728,12 @@ cdef class _Period(PeriodMixin):
         elif util.is_integer_object(other):
             ordinal = self.ordinal + other * self.freq.n
             return Period(ordinal=ordinal, freq=self.freq)
-        elif (PyDateTime_Check(other) or
-              is_period_object(other) or util.is_datetime64_object(other)):
+
+        elif is_period_object(other):
             # can't add datetime-like
-            # GH#17983
+            # GH#17983; can't just return NotImplemented bc we get a RecursionError
+            #  when called via np.add.reduce see TestNumpyReductions.test_add
+            #  in npdev build
             sname = type(self).__name__
             oname = type(other).__name__
             raise TypeError(f"unsupported operand type(s) for +: '{sname}' "
@@ -1750,16 +1752,12 @@ cdef class _Period(PeriodMixin):
                 return NaT
             return NotImplemented
 
-        elif is_any_td_scalar(other):
-            neg_other = -other
-            return self + neg_other
-        elif is_offset_object(other):
-            # Non-Tick DateOffset
-            neg_other = -other
-            return self + neg_other
-        elif util.is_integer_object(other):
-            ordinal = self.ordinal - other * self.freq.n
-            return Period(ordinal=ordinal, freq=self.freq)
+        elif (
+            is_any_td_scalar(other)
+            or is_offset_object(other)
+            or util.is_integer_object(other)
+        ):
+            return self + (-other)
         elif is_period_object(other):
             self._require_matching_freq(other)
             # GH 23915 - mul by base freq since __add__ is agnostic of n

pandas/core/arrays/datetimelike.py

@@ -72,6 +72,7 @@
 from pandas.util._exceptions import find_stack_level
 
 from pandas.core.dtypes.common import (
+    DT64NS_DTYPE,
     is_all_strings,
     is_categorical_dtype,
     is_datetime64_any_dtype,
@@ -89,7 +90,10 @@
     is_unsigned_integer_dtype,
     pandas_dtype,
 )
-from pandas.core.dtypes.dtypes import ExtensionDtype
+from pandas.core.dtypes.dtypes import (
+    DatetimeTZDtype,
+    ExtensionDtype,
+)
 from pandas.core.dtypes.missing import (
     is_valid_na_for_dtype,
     isna,
@@ -113,6 +117,7 @@
 import pandas.core.common as com
 from pandas.core.construction import (
     array as pd_array,
+    ensure_wrapped_if_datetimelike,
     extract_array,
 )
 from pandas.core.indexers import (
@@ -1082,26 +1087,123 @@ def _cmp_method(self, other, op):
     __divmod__ = make_invalid_op("__divmod__")
     __rdivmod__ = make_invalid_op("__rdivmod__")
 
+    @final
     def _add_datetimelike_scalar(self, other):
         # Overridden by TimedeltaArray
-        raise TypeError(f"cannot add {type(self).__name__} and {type(other).__name__}")
+        if not is_timedelta64_dtype(self.dtype):
+            raise TypeError(
+                f"cannot add {type(self).__name__} and {type(other).__name__}"
+            )
 
-    _add_datetime_arraylike = _add_datetimelike_scalar
+        from pandas.core.arrays import DatetimeArray
 
-    def _sub_datetimelike_scalar(self, other):
-        # Overridden by DatetimeArray
         assert other is not NaT
-        raise TypeError(f"cannot subtract a datelike from a {type(self).__name__}")
+        other = Timestamp(other)
+        if other is NaT:
+            # In this case we specifically interpret NaT as a datetime, not
+            # the timedelta interpretation we would get by returning self + NaT
+            result = self.asi8.view("m8[ms]") + NaT.to_datetime64()
+            return DatetimeArray(result)
+
+        i8 = self.asi8
+        result = checked_add_with_arr(i8, other.value, arr_mask=self._isnan)
+        result = self._maybe_mask_results(result)
+        dtype = DatetimeTZDtype(tz=other.tz) if other.tz else DT64NS_DTYPE
+        return DatetimeArray(result, dtype=dtype, freq=self.freq)
+
+    @final
+    def _add_datetime_arraylike(self, other):
+        if not is_timedelta64_dtype(self.dtype):
+            raise TypeError(
+                f"cannot add {type(self).__name__} and {type(other).__name__}"
+            )
+
+        # At this point we have already checked that other.dtype is datetime64
+        other = ensure_wrapped_if_datetimelike(other)
+        # defer to DatetimeArray.__add__
+        return other + self
+
+    @final
+    def _sub_datetimelike_scalar(self, other: datetime | np.datetime64):
+        if self.dtype.kind != "M":
+            raise TypeError(f"cannot subtract a datelike from a {type(self).__name__}")
+
+        self = cast("DatetimeArray", self)
+        # subtract a datetime from myself, yielding a ndarray[timedelta64[ns]]
+
+        # error: Non-overlapping identity check (left operand type: "Union[datetime,
+        # datetime64]", right operand type: "NaTType")  [comparison-overlap]
+        assert other is not NaT  # type: ignore[comparison-overlap]
+        other = Timestamp(other)
+        # error: Non-overlapping identity check (left operand type: "Timestamp",
+        # right operand type: "NaTType")
+        if other is NaT:  # type: ignore[comparison-overlap]
+            return self - NaT
 
-    _sub_datetime_arraylike = _sub_datetimelike_scalar
+        try:
+            self._assert_tzawareness_compat(other)
+        except TypeError as err:
+            new_message = str(err).replace("compare", "subtract")
+            raise type(err)(new_message) from err
+
+        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]")
+
+    @final
+    def _sub_datetime_arraylike(self, other):
+        if self.dtype.kind != "M":
+            raise TypeError(f"cannot subtract a datelike from a {type(self).__name__}")
+
+        if len(self) != len(other):
+            raise ValueError("cannot add indices of unequal length")
 
+        self = cast("DatetimeArray", self)
+        other = ensure_wrapped_if_datetimelike(other)
+
+        try:
+            self._assert_tzawareness_compat(other)
+        except TypeError as err:
+            new_message = str(err).replace("compare", "subtract")
+            raise type(err)(new_message) from err
+
+        self_i8 = self.asi8
+        other_i8 = other.asi8
+        arr_mask = self._isnan | other._isnan
+        new_values = checked_add_with_arr(self_i8, -other_i8, arr_mask=arr_mask)
+        if self._hasna or other._hasna:
+            np.putmask(new_values, arr_mask, iNaT)
+        return new_values.view("timedelta64[ns]")
+
+    @final
     def _sub_period(self, other: Period):
-        # Overridden by PeriodArray
-        raise TypeError(f"cannot subtract Period from a {type(self).__name__}")
+        if not is_period_dtype(self.dtype):
+            raise TypeError(f"cannot subtract Period from a {type(self).__name__}")
+
+        # If the operation is well-defined, we return an object-dtype ndarray
+        # of DateOffsets.  Null entries are filled with pd.NaT
+        self._check_compatible_with(other)
+        asi8 = self.asi8
+        new_data = asi8 - other.ordinal
+        new_data = np.array([self.freq.base * x for x in new_data])
+
+        if self._hasna:
+            new_data[self._isnan] = NaT
+
+        return new_data
 
+    @final
     def _add_period(self, other: Period):
-        # Overridden by TimedeltaArray
-        raise TypeError(f"cannot add Period to a {type(self).__name__}")
+        if not is_timedelta64_dtype(self.dtype):
+            raise TypeError(f"cannot add Period to a {type(self).__name__}")
+
+        # We will wrap in a PeriodArray and defer to the reversed operation
+        from pandas.core.arrays.period import PeriodArray
+
+        i8vals = np.broadcast_to(other.ordinal, self.shape)
+        parr = PeriodArray(i8vals, freq=other.freq)
+        return parr + self
 
     def _add_offset(self, offset):
         raise AbstractMethodError(self)
@@ -1116,9 +1218,9 @@ def _add_timedeltalike_scalar(self, other):
         """
         if isna(other):
             # i.e np.timedelta64("NaT"), not recognized by delta_to_nanoseconds
-            new_values = np.empty(self.shape, dtype="i8")
+            new_values = np.empty(self.shape, dtype="i8").view(self._ndarray.dtype)
             new_values.fill(iNaT)
-            return type(self)(new_values, dtype=self.dtype)
+            return type(self)._simple_new(new_values, dtype=self.dtype)
 
         # PeriodArray overrides, so we only get here with DTA/TDA
         # error: "DatetimeLikeArrayMixin" has no attribute "_reso"
@@ -1139,7 +1241,9 @@ def _add_timedeltalike_scalar(self, other):
             new_values, dtype=self.dtype, freq=new_freq
         )
 
-    def _add_timedelta_arraylike(self, other):
+    def _add_timedelta_arraylike(
+        self, other: TimedeltaArray | npt.NDArray[np.timedelta64]
+    ):
         """
         Add a delta of a TimedeltaIndex
 
@@ -1152,11 +1256,8 @@ def _add_timedelta_arraylike(self, other):
         if len(self) != len(other):
             raise ValueError("cannot add indices of unequal length")
 
-        if isinstance(other, np.ndarray):
-            # ndarray[timedelta64]; wrap in TimedeltaIndex for op
-            from pandas.core.arrays import TimedeltaArray
-
-            other = TimedeltaArray._from_sequence(other)
+        other = ensure_wrapped_if_datetimelike(other)
+        other = cast("TimedeltaArray", other)
 
         self_i8 = self.asi8
         other_i8 = other.asi8
@@ -1200,12 +1301,27 @@ def _sub_nat(self):
         result.fill(iNaT)
         return result.view("timedelta64[ns]")
 
-    def _sub_period_array(self, other):
-        # Overridden by PeriodArray
-        raise TypeError(
-            f"cannot subtract {other.dtype}-dtype from {type(self).__name__}"
+    @final
+    def _sub_period_array(self, other: PeriodArray) -> npt.NDArray[np.object_]:
+        if not is_period_dtype(self.dtype):
+            raise TypeError(
+                f"cannot subtract {other.dtype}-dtype from {type(self).__name__}"
+            )
+
+        self = cast("PeriodArray", self)
+        self._require_matching_freq(other)
+
+        new_values = checked_add_with_arr(
+            self.asi8, -other.asi8, arr_mask=self._isnan, b_mask=other._isnan
         )
 
+        new_values = np.array([self.freq.base * x for x in new_values])
+        if self._hasna or other._hasna:
+            mask = self._isnan | other._isnan
+            new_values[mask] = NaT
+        return new_values
+
+    @final
     def _addsub_object_array(self, other: np.ndarray, op):
         """
         Add or subtract array-like of DateOffset objects

pandas/core/arrays/datetimes.py

@@ -70,7 +70,6 @@
 from pandas.core.dtypes.generic import ABCMultiIndex
 from pandas.core.dtypes.missing import isna
 
-from pandas.core.algorithms import checked_add_with_arr
 from pandas.core.arrays import (
     ExtensionArray,
     datetimelike as dtl,
@@ -733,43 +732,17 @@ def _assert_tzawareness_compat(self, other) -> None:
     # -----------------------------------------------------------------
     # Arithmetic Methods
 
-    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)
-
-        try:
-            self._assert_tzawareness_compat(other)
-        except TypeError as error:
-            new_message = str(error).replace("compare", "subtract")
-            raise type(error)(new_message) from error
-
-        self_i8 = self.asi8
-        other_i8 = other.asi8
-        arr_mask = self._isnan | other._isnan
-        new_values = checked_add_with_arr(self_i8, -other_i8, arr_mask=arr_mask)
-        if self._hasna or other._hasna:
-            np.putmask(new_values, arr_mask, iNaT)
-        return new_values.view("timedelta64[ns]")
-
     def _add_offset(self, offset) -> DatetimeArray:
 
         assert not isinstance(offset, Tick)
+
+        if self.tz is not None:
+            values = self.tz_localize(None)
+        else:
+            values = self
+
         try:
-            if self.tz is not None:
-                values = self.tz_localize(None)
-            else:
-                values = self
             result = offset._apply_array(values).view(values.dtype)
-            result = DatetimeArray._simple_new(result, dtype=result.dtype)
-            if self.tz is not None:
-                # FIXME: tz_localize with non-nano
-                result = result.tz_localize(self.tz)
-
         except NotImplementedError:
             warnings.warn(
                 "Non-vectorized DateOffset being applied to Series or DatetimeIndex.",
@@ -781,35 +754,18 @@ def _add_offset(self, offset) -> DatetimeArray:
                 # GH#30336 _from_sequence won't be able to infer self.tz
                 return result.tz_localize(self.tz)
 
-        return result
-
-    def _sub_datetimelike_scalar(self, other):
-        # subtract a datetime from myself, yielding a ndarray[timedelta64[ns]]
-        assert isinstance(other, (datetime, np.datetime64))
-        # error: Non-overlapping identity check (left operand type: "Union[datetime,
-        # datetime64]", right operand type: "NaTType")  [comparison-overlap]
-        assert other is not NaT  # type: ignore[comparison-overlap]
-        other = Timestamp(other)
-        # error: Non-overlapping identity check (left operand type: "Timestamp",
-        # right operand type: "NaTType")
-        if other is NaT:  # type: ignore[comparison-overlap]
-            return self - NaT
-
-        try:
-            self._assert_tzawareness_compat(other)
-        except TypeError as error:
-            new_message = str(error).replace("compare", "subtract")
-            raise type(error)(new_message) from error
+        else:
+            result = DatetimeArray._simple_new(result, dtype=result.dtype)
+            if self.tz is not None:
+                # FIXME: tz_localize with non-nano
+                result = result.tz_localize(self.tz)
 
-        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]")
+        return result
 
     # -----------------------------------------------------------------
     # Timezone Conversion and Localization Methods
 
-    def _local_timestamps(self) -> np.ndarray:
+    def _local_timestamps(self) -> npt.NDArray[np.int64]:
         """
         Convert to an i8 (unix-like nanosecond timestamp) representation
         while keeping the local timezone and not using UTC.
@@ -1238,7 +1194,7 @@ def to_perioddelta(self, freq) -> TimedeltaArray:
     # -----------------------------------------------------------------
     # Properties - Vectorized Timestamp Properties/Methods
 
-    def month_name(self, locale=None):
+    def month_name(self, locale=None) -> npt.NDArray[np.object_]:
         """
         Return the month names of the :class:`~pandas.Series` or
         :class:`~pandas.DatetimeIndex` with specified locale.
@@ -1283,7 +1239,7 @@ def month_name(self, locale=None):
         result = self._maybe_mask_results(result, fill_value=None)
         return result
 
-    def day_name(self, locale=None):
+    def day_name(self, locale=None) -> npt.NDArray[np.object_]:
         """
         Return the day names of the :class:`~pandas.Series` or
         :class:`~pandas.DatetimeIndex` with specified locale.
@@ -1949,7 +1905,7 @@ def weekofyear(self):
         """,
     )
 
-    def to_julian_date(self) -> np.ndarray:
+    def to_julian_date(self) -> npt.NDArray[np.float64]:
         """
         Convert Datetime Array to float64 ndarray of Julian Dates.
         0 Julian date is noon January 1, 4713 BC.