Skip to content

Commit 855ed4e

Browse files
authored
REF: de-duplicate reso-casting code (#48953)
1 parent ca2d5ad commit 855ed4e

File tree

8 files changed

+44
-43
lines changed

8 files changed

+44
-43
lines changed

pandas/_libs/tslibs/__init__.py

-2
Original file line numberDiff line numberDiff line change
@@ -31,15 +31,13 @@
3131
"periods_per_day",
3232
"periods_per_second",
3333
"is_supported_unit",
34-
"npy_unit_to_abbrev",
3534
]
3635

3736
from pandas._libs.tslibs import dtypes
3837
from pandas._libs.tslibs.conversion import localize_pydatetime
3938
from pandas._libs.tslibs.dtypes import (
4039
Resolution,
4140
is_supported_unit,
42-
npy_unit_to_abbrev,
4341
periods_per_day,
4442
periods_per_second,
4543
)

pandas/_libs/tslibs/timedeltas.pyi

+3
Original file line numberDiff line numberDiff line change
@@ -78,6 +78,7 @@ def delta_to_nanoseconds(
7878
) -> int: ...
7979

8080
class Timedelta(timedelta):
81+
_reso: int
8182
min: ClassVar[Timedelta]
8283
max: ClassVar[Timedelta]
8384
resolution: ClassVar[Timedelta]
@@ -153,4 +154,6 @@ class Timedelta(timedelta):
153154
def freq(self) -> None: ...
154155
@property
155156
def is_populated(self) -> bool: ...
157+
@property
158+
def _unit(self) -> str: ...
156159
def _as_unit(self, unit: str, round_ok: bool = ...) -> Timedelta: ...

pandas/_libs/tslibs/timedeltas.pyx

+7
Original file line numberDiff line numberDiff line change
@@ -1015,6 +1015,13 @@ cdef class _Timedelta(timedelta):
10151015
max = MinMaxReso("max")
10161016
resolution = MinMaxReso("resolution")
10171017

1018+
@property
1019+
def _unit(self) -> str:
1020+
"""
1021+
The abbreviation associated with self._reso.
1022+
"""
1023+
return npy_unit_to_abbrev(self._reso)
1024+
10181025
@property
10191026
def days(self) -> int: # TODO(cython3): make cdef property
10201027
# NB: using the python C-API PyDateTime_DELTA_GET_DAYS will fail

pandas/_libs/tslibs/timestamps.pyi

+2
Original file line numberDiff line numberDiff line change
@@ -222,4 +222,6 @@ class Timestamp(datetime):
222222
def days_in_month(self) -> int: ...
223223
@property
224224
def daysinmonth(self) -> int: ...
225+
@property
226+
def _unit(self) -> str: ...
225227
def _as_unit(self, unit: str, round_ok: bool = ...) -> Timestamp: ...

pandas/_libs/tslibs/timestamps.pyx

+7
Original file line numberDiff line numberDiff line change
@@ -257,6 +257,13 @@ cdef class _Timestamp(ABCTimestamp):
257257
)
258258
return self._freq
259259

260+
@property
261+
def _unit(self) -> str:
262+
"""
263+
The abbreviation associated with self._reso.
264+
"""
265+
return npy_unit_to_abbrev(self._reso)
266+
260267
# -----------------------------------------------------------------
261268
# Constructors
262269

pandas/core/arrays/datetimelike.py

+20-32
Original file line numberDiff line numberDiff line change
@@ -44,7 +44,6 @@
4444
iNaT,
4545
ints_to_pydatetime,
4646
ints_to_pytimedelta,
47-
npy_unit_to_abbrev,
4847
to_offset,
4948
)
5049
from pandas._libs.tslibs.fields import (
@@ -1169,13 +1168,8 @@ def _add_datetimelike_scalar(self, other) -> DatetimeArray:
11691168
# Preserve our resolution
11701169
return DatetimeArray._simple_new(result, dtype=result.dtype)
11711170

1172-
if self._reso != other._reso:
1173-
# Just as with Timestamp/Timedelta, we cast to the higher resolution
1174-
if self._reso < other._reso:
1175-
unit = npy_unit_to_abbrev(other._reso)
1176-
self = self._as_unit(unit)
1177-
else:
1178-
other = other._as_unit(self._unit)
1171+
self, other = self._ensure_matching_resos(other)
1172+
self = cast("TimedeltaArray", self)
11791173

11801174
i8 = self.asi8
11811175
result = checked_add_with_arr(i8, other.value, arr_mask=self._isnan)
@@ -1208,16 +1202,10 @@ def _sub_datetimelike_scalar(self, other: datetime | np.datetime64):
12081202
# i.e. np.datetime64("NaT")
12091203
return self - NaT
12101204

1211-
other = Timestamp(other)
1205+
ts = Timestamp(other)
12121206

1213-
if other._reso != self._reso:
1214-
if other._reso < self._reso:
1215-
other = other._as_unit(self._unit)
1216-
else:
1217-
unit = npy_unit_to_abbrev(other._reso)
1218-
self = self._as_unit(unit)
1219-
1220-
return self._sub_datetimelike(other)
1207+
self, ts = self._ensure_matching_resos(ts)
1208+
return self._sub_datetimelike(ts)
12211209

12221210
@final
12231211
def _sub_datetime_arraylike(self, other):
@@ -1230,12 +1218,7 @@ def _sub_datetime_arraylike(self, other):
12301218
self = cast("DatetimeArray", self)
12311219
other = ensure_wrapped_if_datetimelike(other)
12321220

1233-
if other._reso != self._reso:
1234-
if other._reso < self._reso:
1235-
other = other._as_unit(self._unit)
1236-
else:
1237-
self = self._as_unit(other._unit)
1238-
1221+
self, other = self._ensure_matching_resos(other)
12391222
return self._sub_datetimelike(other)
12401223

12411224
@final
@@ -1319,17 +1302,11 @@ def _add_timedelta_arraylike(
13191302
raise ValueError("cannot add indices of unequal length")
13201303

13211304
other = ensure_wrapped_if_datetimelike(other)
1322-
other = cast("TimedeltaArray", other)
1305+
tda = cast("TimedeltaArray", other)
13231306
self = cast("DatetimeArray | TimedeltaArray", self)
13241307

1325-
if self._reso != other._reso:
1326-
# Just as with Timestamp/Timedelta, we cast to the higher resolution
1327-
if self._reso < other._reso:
1328-
self = self._as_unit(other._unit)
1329-
else:
1330-
other = other._as_unit(self._unit)
1331-
1332-
return self._add_timedeltalike(other)
1308+
self, tda = self._ensure_matching_resos(tda)
1309+
return self._add_timedeltalike(tda)
13331310

13341311
@final
13351312
def _add_timedeltalike(self, other: Timedelta | TimedeltaArray):
@@ -2098,6 +2075,17 @@ def _as_unit(self: TimelikeOpsT, unit: str) -> TimelikeOpsT:
20982075
new_values, dtype=new_dtype, freq=self.freq # type: ignore[call-arg]
20992076
)
21002077

2078+
# TODO: annotate other as DatetimeArray | TimedeltaArray | Timestamp | Timedelta
2079+
# with the return type matching input type. TypeVar?
2080+
def _ensure_matching_resos(self, other):
2081+
if self._reso != other._reso:
2082+
# Just as with Timestamp/Timedelta, we cast to the higher resolution
2083+
if self._reso < other._reso:
2084+
self = self._as_unit(other._unit)
2085+
else:
2086+
other = other._as_unit(self._unit)
2087+
return self, other
2088+
21012089
# --------------------------------------------------------------
21022090

21032091
def __array_ufunc__(self, ufunc: np.ufunc, method: str, *inputs, **kwargs):

pandas/tests/scalar/timestamp/test_timestamp.py

+5-8
Original file line numberDiff line numberDiff line change
@@ -18,10 +18,7 @@
1818
utc,
1919
)
2020

21-
from pandas._libs.tslibs.dtypes import (
22-
NpyDatetimeUnit,
23-
npy_unit_to_abbrev,
24-
)
21+
from pandas._libs.tslibs.dtypes import NpyDatetimeUnit
2522
from pandas._libs.tslibs.timezones import (
2623
dateutil_gettz as gettz,
2724
get_timezone,
@@ -964,7 +961,7 @@ def test_sub_datetimelike_mismatched_reso(self, ts_tz):
964961
if ts._reso < other._reso:
965962
# Case where rounding is lossy
966963
other2 = other + Timedelta._from_value_and_reso(1, other._reso)
967-
exp = ts._as_unit(npy_unit_to_abbrev(other._reso)) - other2
964+
exp = ts._as_unit(other._unit) - other2
968965

969966
res = ts - other2
970967
assert res == exp
@@ -975,7 +972,7 @@ def test_sub_datetimelike_mismatched_reso(self, ts_tz):
975972
assert res._reso == max(ts._reso, other._reso)
976973
else:
977974
ts2 = ts + Timedelta._from_value_and_reso(1, ts._reso)
978-
exp = ts2 - other._as_unit(npy_unit_to_abbrev(ts2._reso))
975+
exp = ts2 - other._as_unit(ts2._unit)
979976

980977
res = ts2 - other
981978
assert res == exp
@@ -1012,7 +1009,7 @@ def test_sub_timedeltalike_mismatched_reso(self, ts_tz):
10121009
if ts._reso < other._reso:
10131010
# Case where rounding is lossy
10141011
other2 = other + Timedelta._from_value_and_reso(1, other._reso)
1015-
exp = ts._as_unit(npy_unit_to_abbrev(other._reso)) + other2
1012+
exp = ts._as_unit(other._unit) + other2
10161013
res = ts + other2
10171014
assert res == exp
10181015
assert res._reso == max(ts._reso, other._reso)
@@ -1021,7 +1018,7 @@ def test_sub_timedeltalike_mismatched_reso(self, ts_tz):
10211018
assert res._reso == max(ts._reso, other._reso)
10221019
else:
10231020
ts2 = ts + Timedelta._from_value_and_reso(1, ts._reso)
1024-
exp = ts2 + other._as_unit(npy_unit_to_abbrev(ts2._reso))
1021+
exp = ts2 + other._as_unit(ts2._unit)
10251022

10261023
res = ts2 + other
10271024
assert res == exp

pandas/tests/tslibs/test_api.py

-1
Original file line numberDiff line numberDiff line change
@@ -56,7 +56,6 @@ def test_namespace():
5656
"periods_per_day",
5757
"periods_per_second",
5858
"is_supported_unit",
59-
"npy_unit_to_abbrev",
6059
]
6160

6261
expected = set(submodules + api)

0 commit comments

Comments
 (0)