diff --git a/pandas/core/arrays/datetimelike.py b/pandas/core/arrays/datetimelike.py index 997cef7b09576..f29930b010795 100644 --- a/pandas/core/arrays/datetimelike.py +++ b/pandas/core/arrays/datetimelike.py @@ -1112,6 +1112,41 @@ def _cmp_method(self, other, op): __divmod__ = make_invalid_op("__divmod__") __rdivmod__ = make_invalid_op("__rdivmod__") + @final + def _get_i8_values_and_mask( + self, other + ) -> tuple[int | npt.NDArray[np.int64], None | npt.NDArray[np.bool_]]: + """ + Get the int64 values and b_mask to pass to checked_add_with_arr. + """ + if isinstance(other, Period): + i8values = other.ordinal + mask = None + elif isinstance(other, Timestamp): + i8values = other.value + mask = None + else: + # PeriodArray, DatetimeArray, TimedeltaArray + mask = other._isnan + i8values = other.asi8 + return i8values, mask + + @final + def _get_arithmetic_result_freq(self, other) -> BaseOffset | None: + """ + Check if we can preserve self.freq in addition or subtraction. + """ + # Adding or subtracting a Timedelta/Timestamp scalar is freq-preserving + # whenever self.freq is a Tick + if is_period_dtype(self.dtype): + return self.freq + elif not lib.is_scalar(other): + return None + elif isinstance(self.freq, Tick): + # In these cases + return self.freq + return None + @final def _add_datetimelike_scalar(self, other) -> DatetimeArray: if not is_timedelta64_dtype(self.dtype): @@ -1228,15 +1263,7 @@ def _sub_period(self, other: Period) -> npt.NDArray[np.object_]: # 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) - new_i8_data = checked_add_with_arr( - self.asi8, -other.ordinal, arr_mask=self._isnan - ) - new_data = np.array([self.freq.base * x for x in new_i8_data]) - - if self._hasna: - new_data[self._isnan] = NaT - - return new_data + return self._sub_periodlike(other) @final def _add_period(self, other: Period) -> PeriodArray: @@ -1361,15 +1388,26 @@ def _sub_period_array(self, other: PeriodArray) -> npt.NDArray[np.object_]: self = cast("PeriodArray", self) self._require_matching_freq(other) - new_i8_values = checked_add_with_arr( - self.asi8, -other.asi8, arr_mask=self._isnan, b_mask=other._isnan + return self._sub_periodlike(other) + + @final + def _sub_periodlike(self, other: Period | PeriodArray) -> npt.NDArray[np.object_]: + # caller is responsible for calling + # require_matching_freq/check_compatible_with + other_i8, o_mask = self._get_i8_values_and_mask(other) + new_i8_data = checked_add_with_arr( + self.asi8, -other_i8, arr_mask=self._isnan, b_mask=o_mask ) + new_data = np.array([self.freq.base * x for x in new_i8_data]) - new_values = np.array([self.freq.base * x for x in new_i8_values]) - if self._hasna or other._hasna: - mask = self._isnan | other._isnan - new_values[mask] = NaT - return new_values + if o_mask is None: + # i.e. Period scalar + mask = self._isnan + else: + # i.e. PeriodArray + mask = self._isnan | o_mask + new_data[mask] = NaT + return new_data @final def _addsub_object_array(self, other: np.ndarray, op):