diff --git a/pandas/core/algorithms.py b/pandas/core/algorithms.py index 888e943488953..db76f5c51752f 100644 --- a/pandas/core/algorithms.py +++ b/pandas/core/algorithms.py @@ -1017,10 +1017,10 @@ def rank( def checked_add_with_arr( arr: npt.NDArray[np.int64], - b, + b: int | npt.NDArray[np.int64], arr_mask: npt.NDArray[np.bool_] | None = None, b_mask: npt.NDArray[np.bool_] | None = None, -) -> np.ndarray: +) -> npt.NDArray[np.int64]: """ Perform array addition that checks for underflow and overflow. @@ -1098,7 +1098,12 @@ def checked_add_with_arr( if to_raise: raise OverflowError("Overflow in int64 addition") - return arr + b + + result = arr + b + if arr_mask is not None or b2_mask is not None: + np.putmask(result, ~not_nan, iNaT) + + return result # --------------- # diff --git a/pandas/core/arrays/datetimelike.py b/pandas/core/arrays/datetimelike.py index 1dfb070e29c30..122b61e2e351d 100644 --- a/pandas/core/arrays/datetimelike.py +++ b/pandas/core/arrays/datetimelike.py @@ -1104,8 +1104,12 @@ def _add_datetimelike_scalar(self, other): return DatetimeArray(result) i8 = self.asi8 - result = checked_add_with_arr(i8, other.value, arr_mask=self._isnan) - result = self._maybe_mask_results(result) + # Incompatible types in assignment (expression has type "ndarray[Any, + # dtype[signedinteger[_64Bit]]]", variable has type + # "ndarray[Any, dtype[datetime64]]") + result = checked_add_with_arr( # type: ignore[assignment] + i8, other.value, arr_mask=self._isnan + ) dtype = DatetimeTZDtype(tz=other.tz) if other.tz else DT64NS_DTYPE return DatetimeArray(result, dtype=dtype, freq=self.freq) @@ -1146,7 +1150,6 @@ def _sub_datetimelike_scalar(self, other: datetime | np.datetime64): 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 @@ -1168,14 +1171,13 @@ def _sub_datetime_arraylike(self, other): 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) + new_values = checked_add_with_arr( + self_i8, -other_i8, arr_mask=self._isnan, b_mask=other._isnan + ) return new_values.view("timedelta64[ns]") @final - def _sub_period(self, other: Period): + def _sub_period(self, other: Period) -> npt.NDArray[np.object_]: if not is_period_dtype(self.dtype): raise TypeError(f"cannot subtract Period from a {type(self).__name__}") @@ -1183,8 +1185,8 @@ def _sub_period(self, other: Period): # 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]) + new_i8_data = asi8 - other.ordinal # TODO: checked_add_with_arr + new_data = np.array([self.freq.base * x for x in new_i8_data]) if self._hasna: new_data[self._isnan] = NaT @@ -1192,7 +1194,7 @@ def _sub_period(self, other: Period): return new_data @final - def _add_period(self, other: Period): + def _add_period(self, other: Period) -> PeriodArray: if not is_timedelta64_dtype(self.dtype): raise TypeError(f"cannot add Period to a {type(self).__name__}") @@ -1225,8 +1227,6 @@ def _add_timedeltalike_scalar(self, other): inc = delta_to_nanoseconds(other, reso=self._reso) # type: ignore[attr-defined] new_values = checked_add_with_arr(self.asi8, inc, arr_mask=self._isnan) - new_values = new_values.view("i8") - new_values = self._maybe_mask_results(new_values) new_values = new_values.view(self._ndarray.dtype) new_freq = None @@ -1262,10 +1262,6 @@ def _add_timedelta_arraylike( new_values = checked_add_with_arr( self_i8, other_i8, arr_mask=self._isnan, b_mask=other._isnan ) - if self._hasna or other._hasna: - mask = self._isnan | other._isnan - np.putmask(new_values, mask, iNaT) - return type(self)(new_values, dtype=self.dtype) @final @@ -1309,11 +1305,11 @@ def _sub_period_array(self, other: PeriodArray) -> npt.NDArray[np.object_]: self = cast("PeriodArray", self) self._require_matching_freq(other) - new_values = checked_add_with_arr( + new_i8_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]) + 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 diff --git a/pandas/core/arrays/period.py b/pandas/core/arrays/period.py index b6d21cd9dac54..cab9a8a565145 100644 --- a/pandas/core/arrays/period.py +++ b/pandas/core/arrays/period.py @@ -733,8 +733,6 @@ def _addsub_int_array_or_scalar( if op is operator.sub: other = -other res_values = algos.checked_add_with_arr(self.asi8, other, arr_mask=self._isnan) - res_values = res_values.view("i8") - np.putmask(res_values, self._isnan, iNaT) return type(self)(res_values, freq=self.freq) def _add_offset(self, other: BaseOffset):