Skip to content

Commit 8f26ab7

Browse files
authored
REF: helpers to de-duplicate datetimelike arithmetic (#48844)
1 parent c7623b3 commit 8f26ab7

File tree

1 file changed

+54
-16
lines changed

1 file changed

+54
-16
lines changed

pandas/core/arrays/datetimelike.py

+54-16
Original file line numberDiff line numberDiff line change
@@ -1112,6 +1112,41 @@ def _cmp_method(self, other, op):
11121112
__divmod__ = make_invalid_op("__divmod__")
11131113
__rdivmod__ = make_invalid_op("__rdivmod__")
11141114

1115+
@final
1116+
def _get_i8_values_and_mask(
1117+
self, other
1118+
) -> tuple[int | npt.NDArray[np.int64], None | npt.NDArray[np.bool_]]:
1119+
"""
1120+
Get the int64 values and b_mask to pass to checked_add_with_arr.
1121+
"""
1122+
if isinstance(other, Period):
1123+
i8values = other.ordinal
1124+
mask = None
1125+
elif isinstance(other, Timestamp):
1126+
i8values = other.value
1127+
mask = None
1128+
else:
1129+
# PeriodArray, DatetimeArray, TimedeltaArray
1130+
mask = other._isnan
1131+
i8values = other.asi8
1132+
return i8values, mask
1133+
1134+
@final
1135+
def _get_arithmetic_result_freq(self, other) -> BaseOffset | None:
1136+
"""
1137+
Check if we can preserve self.freq in addition or subtraction.
1138+
"""
1139+
# Adding or subtracting a Timedelta/Timestamp scalar is freq-preserving
1140+
# whenever self.freq is a Tick
1141+
if is_period_dtype(self.dtype):
1142+
return self.freq
1143+
elif not lib.is_scalar(other):
1144+
return None
1145+
elif isinstance(self.freq, Tick):
1146+
# In these cases
1147+
return self.freq
1148+
return None
1149+
11151150
@final
11161151
def _add_datetimelike_scalar(self, other) -> DatetimeArray:
11171152
if not is_timedelta64_dtype(self.dtype):
@@ -1228,15 +1263,7 @@ def _sub_period(self, other: Period) -> npt.NDArray[np.object_]:
12281263
# If the operation is well-defined, we return an object-dtype ndarray
12291264
# of DateOffsets. Null entries are filled with pd.NaT
12301265
self._check_compatible_with(other)
1231-
new_i8_data = checked_add_with_arr(
1232-
self.asi8, -other.ordinal, arr_mask=self._isnan
1233-
)
1234-
new_data = np.array([self.freq.base * x for x in new_i8_data])
1235-
1236-
if self._hasna:
1237-
new_data[self._isnan] = NaT
1238-
1239-
return new_data
1266+
return self._sub_periodlike(other)
12401267

12411268
@final
12421269
def _add_period(self, other: Period) -> PeriodArray:
@@ -1361,15 +1388,26 @@ def _sub_period_array(self, other: PeriodArray) -> npt.NDArray[np.object_]:
13611388
self = cast("PeriodArray", self)
13621389
self._require_matching_freq(other)
13631390

1364-
new_i8_values = checked_add_with_arr(
1365-
self.asi8, -other.asi8, arr_mask=self._isnan, b_mask=other._isnan
1391+
return self._sub_periodlike(other)
1392+
1393+
@final
1394+
def _sub_periodlike(self, other: Period | PeriodArray) -> npt.NDArray[np.object_]:
1395+
# caller is responsible for calling
1396+
# require_matching_freq/check_compatible_with
1397+
other_i8, o_mask = self._get_i8_values_and_mask(other)
1398+
new_i8_data = checked_add_with_arr(
1399+
self.asi8, -other_i8, arr_mask=self._isnan, b_mask=o_mask
13661400
)
1401+
new_data = np.array([self.freq.base * x for x in new_i8_data])
13671402

1368-
new_values = np.array([self.freq.base * x for x in new_i8_values])
1369-
if self._hasna or other._hasna:
1370-
mask = self._isnan | other._isnan
1371-
new_values[mask] = NaT
1372-
return new_values
1403+
if o_mask is None:
1404+
# i.e. Period scalar
1405+
mask = self._isnan
1406+
else:
1407+
# i.e. PeriodArray
1408+
mask = self._isnan | o_mask
1409+
new_data[mask] = NaT
1410+
return new_data
13731411

13741412
@final
13751413
def _addsub_object_array(self, other: np.ndarray, op):

0 commit comments

Comments
 (0)