Skip to content

Commit 5df1c51

Browse files
bashtageKevin Sheppard
bashtage
and
Kevin Sheppard
authored
ENH: Improve typing for Timedelta (#388)
* ENH: Improve typing for Timedelta * ENH: Improve typing of timedelta * ENH/TST: Improve Timedelta and its tests * ENH: Final changes for timedelta * ENH: Further improvements to timestamp * ENH: Add more types * Final changes to timedelta * Final cleanup * CLN: Final clean * CLN: Final clean Co-authored-by: Kevin Sheppard <kevin.sheppard@gmail.com>
1 parent 56bcec7 commit 5df1c51

File tree

4 files changed

+995
-53
lines changed

4 files changed

+995
-53
lines changed

‎pandas-stubs/_libs/tslibs/timedeltas.pyi‎

Lines changed: 276 additions & 37 deletions
Original file line numberDiff line numberDiff line change
@@ -1,24 +1,56 @@
1+
import datetime as dt
12
from datetime import timedelta
23
from typing import (
34
ClassVar,
45
Literal,
6+
NamedTuple,
57
TypeVar,
68
Union,
79
overload,
810
)
911

1012
import numpy as np
13+
import pandas as pd
14+
from pandas import (
15+
DatetimeIndex,
16+
Float64Index,
17+
Int64Index,
18+
PeriodIndex,
19+
Series,
20+
TimedeltaIndex,
21+
)
22+
from pandas.core.series import (
23+
TimedeltaSeries,
24+
TimestampSeries,
25+
)
1126
from typing_extensions import TypeAlias
1227

1328
from pandas._libs.tslibs import (
29+
BaseOffset,
1430
NaTType,
15-
Tick,
1631
)
32+
from pandas._libs.tslibs.period import Period
33+
from pandas._libs.tslibs.timestamps import Timestamp
1734
from pandas._typing import npt
1835

36+
class Components(NamedTuple):
37+
days: int
38+
hours: int
39+
minutes: int
40+
seconds: int
41+
milliseconds: int
42+
microseconds: int
43+
nanoseconds: int
44+
1945
# This should be kept consistent with the keys in the dict timedelta_abbrevs
2046
# in pandas/_libs/tslibs/timedeltas.pyx
2147
TimeDeltaUnitChoices: TypeAlias = Literal[
48+
"H",
49+
"T",
50+
"S",
51+
"L",
52+
"U",
53+
"N",
2254
"W",
2355
"w",
2456
"D",
@@ -70,35 +102,32 @@ UnitChoices: TypeAlias = Union[
70102

71103
_S = TypeVar("_S", bound=timedelta)
72104

73-
def ints_to_pytimedelta(
74-
arr: npt.NDArray[np.int64], # const int64_t[:]
75-
box: bool = ...,
76-
) -> npt.NDArray[np.object_]: ...
77-
def array_to_timedelta64(
78-
values: npt.NDArray[np.object_],
79-
unit: str | None = ...,
80-
errors: str = ...,
81-
) -> np.ndarray: ... # np.ndarray[m8ns]
82-
def parse_timedelta_unit(unit: str | None) -> UnitChoices: ...
83-
def delta_to_nanoseconds(delta: np.timedelta64 | timedelta | Tick) -> int: ...
84-
85105
class Timedelta(timedelta):
86106
min: ClassVar[Timedelta]
87107
max: ClassVar[Timedelta]
88108
resolution: ClassVar[Timedelta]
89-
value: int # np.int64
109+
value: int
90110
def __new__(
91111
cls: type[_S],
92-
value=...,
93-
unit: str = ...,
94-
**kwargs: float | np.integer | np.floating,
112+
value: str | int | Timedelta | timedelta | np.timedelta64 = ...,
113+
unit: TimeDeltaUnitChoices = ...,
114+
*,
115+
days: float | np.integer | np.floating = ...,
116+
seconds: float | np.integer | np.floating = ...,
117+
microseconds: float | np.integer | np.floating = ...,
118+
milliseconds: float | np.integer | np.floating = ...,
119+
minutes: float | np.integer | np.floating = ...,
120+
hours: float | np.integer | np.floating = ...,
121+
weeks: float | np.integer | np.floating = ...,
95122
) -> _S: ...
96123
# GH 46171
97124
# While Timedelta can return pd.NaT, having the constructor return
98125
# a Union with NaTType makes things awkward for users of pandas
99126
@property
100127
def days(self) -> int: ...
101128
@property
129+
def nanoseconds(self) -> int: ...
130+
@property
102131
def seconds(self) -> int: ...
103132
@property
104133
def microseconds(self) -> int: ...
@@ -108,49 +137,259 @@ class Timedelta(timedelta):
108137
@property
109138
def asm8(self) -> np.timedelta64: ...
110139
# TODO: round/floor/ceil could return NaT?
111-
def round(self: _S, freq: str) -> _S: ...
112-
def floor(self: _S, freq: str) -> _S: ...
113-
def ceil(self: _S, freq: str) -> _S: ...
140+
def round(self: _S, freq: str | BaseOffset) -> _S: ...
141+
def floor(self: _S, freq: str | BaseOffset) -> _S: ...
142+
def ceil(self: _S, freq: str | BaseOffset) -> _S: ...
114143
@property
115144
def resolution_string(self) -> str: ...
116-
def __add__(self, other: timedelta) -> Timedelta: ...
117-
def __radd__(self, other: timedelta) -> Timedelta: ...
118-
def __sub__(self, other: timedelta) -> Timedelta: ...
119-
def __rsub__(self, other: timedelta) -> Timedelta: ...
145+
# Override due to more types supported than dt.timedelta
146+
@overload # type: ignore[override]
147+
def __add__(self, other: timedelta | Timedelta | np.timedelta64) -> Timedelta: ...
148+
@overload
149+
def __add__(self, other: dt.datetime | np.datetime64 | Timestamp) -> Timestamp: ...
150+
@overload
151+
def __add__(self, other: NaTType) -> NaTType: ...
152+
@overload
153+
def __add__(self, other: Period) -> Period: ...
154+
@overload
155+
def __add__(self, other: dt.date) -> dt.date: ...
156+
@overload
157+
def __add__(self, other: PeriodIndex) -> PeriodIndex: ...
158+
@overload
159+
def __add__(self, other: DatetimeIndex) -> DatetimeIndex: ...
160+
@overload
161+
def __add__(
162+
self, other: npt.NDArray[np.timedelta64]
163+
) -> npt.NDArray[np.timedelta64]: ...
164+
@overload
165+
def __add__(
166+
self, other: npt.NDArray[np.datetime64]
167+
) -> npt.NDArray[np.datetime64]: ...
168+
@overload
169+
def __add__(self, other: pd.TimedeltaIndex) -> pd.TimedeltaIndex: ...
170+
@overload
171+
def __add__(
172+
self, other: TimedeltaSeries | Series[pd.Timedelta]
173+
) -> TimedeltaSeries: ...
174+
@overload
175+
def __add__(
176+
self, other: Series[Timestamp] | TimestampSeries
177+
) -> TimestampSeries: ...
178+
@overload
179+
def __radd__(self, other: np.datetime64) -> Timestamp: ...
180+
@overload
181+
def __radd__(self, other: timedelta | Timedelta | np.timedelta64) -> Timedelta: ...
182+
@overload
183+
def __radd__(self, other: NaTType) -> NaTType: ...
184+
@overload
185+
def __radd__(
186+
self, other: npt.NDArray[np.timedelta64]
187+
) -> npt.NDArray[np.timedelta64]: ...
188+
@overload
189+
def __radd__(
190+
self, other: npt.NDArray[np.datetime64]
191+
) -> npt.NDArray[np.datetime64]: ...
192+
@overload
193+
def __radd__(self, other: pd.TimedeltaIndex) -> pd.TimedeltaIndex: ...
194+
@overload
195+
def __radd__(self, other: pd.PeriodIndex) -> pd.PeriodIndex: ...
196+
# Override due to more types supported than dt.timedelta
197+
@overload # type: ignore[override]
198+
def __sub__(self, other: timedelta | Timedelta | np.timedelta64) -> Timedelta: ...
199+
@overload
200+
def __sub__(self, other: NaTType) -> NaTType: ...
201+
@overload
202+
def __sub__(
203+
self, other: npt.NDArray[np.timedelta64]
204+
) -> npt.NDArray[np.timedelta64]: ...
205+
@overload
206+
def __sub__(self, other: pd.TimedeltaIndex) -> TimedeltaIndex: ...
207+
@overload
208+
def __sub__(
209+
self, other: TimedeltaSeries | Series[pd.Timedelta]
210+
) -> TimedeltaSeries: ...
211+
@overload
212+
def __rsub__(self, other: timedelta | Timedelta | np.timedelta64) -> Timedelta: ...
213+
@overload
214+
def __rsub__(self, other: Timestamp | np.datetime64) -> Timestamp: ...
215+
@overload
216+
def __rsub__(self, other: NaTType) -> NaTType: ...
217+
@overload
218+
def __rsub__(self, other: Period) -> Period: ...
219+
@overload
220+
def __rsub__(self, other: PeriodIndex) -> PeriodIndex: ...
221+
@overload
222+
def __rsub__(self, other: DatetimeIndex) -> DatetimeIndex: ...
223+
@overload
224+
def __rsub__(
225+
self, other: npt.NDArray[np.datetime64]
226+
) -> npt.NDArray[np.datetime64]: ...
227+
@overload
228+
def __rsub__(
229+
self, other: npt.NDArray[np.timedelta64]
230+
) -> npt.NDArray[np.timedelta64]: ...
231+
@overload
232+
def __rsub__(self, other: pd.TimedeltaIndex) -> pd.TimedeltaIndex: ...
120233
def __neg__(self) -> Timedelta: ...
121234
def __pos__(self) -> Timedelta: ...
122235
def __abs__(self) -> Timedelta: ...
236+
# Override due to more types supported than dt.timedelta
237+
@overload # type: ignore[override]
123238
def __mul__(self, other: float) -> Timedelta: ...
239+
@overload
240+
def __mul__(
241+
self, other: npt.NDArray[np.integer] | npt.NDArray[np.floating]
242+
) -> npt.NDArray[np.timedelta64]: ...
243+
@overload
244+
def __mul__(self, other: Series[int]) -> TimedeltaSeries: ...
245+
@overload
246+
def __mul__(self, other: Series[float]) -> TimedeltaSeries: ...
247+
@overload
248+
def __mul__(self, other: Int64Index | Float64Index) -> TimedeltaIndex: ...
249+
@overload
124250
def __rmul__(self, other: float) -> Timedelta: ...
251+
@overload
252+
def __rmul__(self, other: np.ndarray) -> np.ndarray: ...
253+
@overload
254+
def __rmul__(self, other: Series[int]) -> TimedeltaSeries: ...
255+
@overload
256+
def __rmul__(self, other: Series[float]) -> TimedeltaSeries: ...
257+
@overload
258+
def __rmul__(self, other: Int64Index | Float64Index) -> TimedeltaIndex: ...
259+
# Override due to more types supported than dt.timedelta
125260
# error: Signature of "__floordiv__" incompatible with supertype "timedelta"
126261
@overload # type: ignore[override]
127-
def __floordiv__(self, other: timedelta) -> int: ...
262+
def __floordiv__(self, other: timedelta | Timedelta | np.timedelta64) -> int: ...
128263
@overload
129264
def __floordiv__(self, other: float) -> Timedelta: ...
130265
@overload
266+
def __floordiv__(
267+
self, other: npt.NDArray[np.integer] | npt.NDArray[np.floating]
268+
) -> npt.NDArray[np.timedelta64]: ...
269+
@overload
131270
def __floordiv__(
132271
self, other: npt.NDArray[np.timedelta64]
133-
) -> npt.NDArray[np.intp]: ...
272+
) -> npt.NDArray[np.int_]: ...
273+
@overload
274+
def __floordiv__(self, other: Int64Index | Float64Index) -> TimedeltaIndex: ...
275+
@overload
276+
def __floordiv__(self, other: Series[int]) -> TimedeltaSeries: ...
277+
@overload
278+
def __floordiv__(self, other: Series[float]) -> TimedeltaSeries: ...
134279
@overload
135280
def __floordiv__(
136-
self, other: npt.NDArray[np.number]
137-
) -> npt.NDArray[np.timedelta64] | Timedelta: ...
281+
self, other: Series[Timedelta] | TimedeltaSeries
282+
) -> Series[int]: ...
138283
@overload
139-
def __rfloordiv__(self, other: timedelta | str) -> int: ...
284+
def __floordiv__(self, other: NaTType | None) -> float: ...
140285
@overload
141-
def __rfloordiv__(self, other: NaTType | None) -> NaTType: ...
286+
def __rfloordiv__(self, other: timedelta | Timedelta | str) -> int: ...
142287
@overload
143-
def __rfloordiv__(self, other: np.ndarray) -> npt.NDArray[np.timedelta64]: ...
288+
def __rfloordiv__(self, other: NaTType | None) -> float: ...
144289
@overload
145-
def __truediv__(self, other: timedelta) -> float: ...
290+
def __rfloordiv__(
291+
self, other: npt.NDArray[np.timedelta64]
292+
) -> npt.NDArray[np.int_]: ...
293+
# Override due to more types supported than dt.timedelta
294+
@overload # type: ignore[override]
295+
def __truediv__(self, other: timedelta | Timedelta | NaTType) -> float: ...
146296
@overload
147297
def __truediv__(self, other: float) -> Timedelta: ...
298+
@overload
299+
def __truediv__(
300+
self, other: npt.NDArray[np.integer] | npt.NDArray[np.floating]
301+
) -> npt.NDArray[np.timedelta64]: ...
302+
@overload
303+
def __truediv__(
304+
self, other: Series[Timedelta] | TimedeltaSeries
305+
) -> Series[float]: ...
306+
@overload
307+
def __truediv__(self, other: Series[int]) -> TimedeltaSeries: ...
308+
@overload
309+
def __truediv__(self, other: Series[float]) -> TimedeltaSeries: ...
310+
@overload
311+
def __truediv__(self, other: Int64Index | Float64Index) -> TimedeltaIndex: ...
312+
def __rtruediv__(self, other: timedelta | Timedelta | NaTType) -> float: ...
313+
# Override due to more types supported than dt.timedelta
314+
@overload
315+
def __eq__(self, other: timedelta | Timedelta | np.timedelta64) -> bool: ... # type: ignore[misc]
316+
@overload
317+
def __eq__(self, other: TimedeltaSeries | Series[pd.Timedelta]) -> Series[bool]: ... # type: ignore[misc]
318+
@overload
319+
def __eq__( # type: ignore[misc]
320+
self, other: TimedeltaIndex | npt.NDArray[np.timedelta64]
321+
) -> npt.NDArray[np.bool_]: ...
322+
@overload
323+
def __eq__(self, other: object) -> Literal[False]: ...
324+
# Override due to more types supported than dt.timedelta
325+
@overload
326+
def __ne__(self, other: timedelta | Timedelta | np.timedelta64) -> bool: ... # type: ignore[misc]
327+
@overload
328+
def __ne__(self, other: TimedeltaSeries | Series[pd.Timedelta]) -> Series[bool]: ... # type: ignore[misc]
329+
@overload
330+
def __ne__( # type: ignore[misc]
331+
self, other: TimedeltaIndex | npt.NDArray[np.timedelta64]
332+
) -> npt.NDArray[np.bool_]: ...
333+
@overload
334+
def __ne__(self, other: object) -> Literal[True]: ...
335+
# Override due to more types supported than dt.timedelta
336+
@overload # type: ignore[override]
148337
def __mod__(self, other: timedelta) -> Timedelta: ...
338+
@overload
339+
def __mod__(self, other: float) -> Timedelta: ...
340+
@overload
341+
def __mod__(self, other: Series[int] | Series[float]) -> TimedeltaSeries: ...
342+
@overload
343+
def __mod__(self, other: Int64Index | Float64Index) -> TimedeltaIndex: ...
344+
@overload
345+
def __mod__(
346+
self, other: npt.NDArray[np.integer] | npt.NDArray[np.floating]
347+
) -> npt.NDArray[np.timedelta64]: ...
348+
@overload
349+
def __mod__(
350+
self, other: Series[int] | Series[float] | Series[Timedelta] | TimedeltaSeries
351+
) -> TimedeltaSeries: ...
149352
def __divmod__(self, other: timedelta) -> tuple[int, Timedelta]: ...
150-
def __le__(self, other: timedelta) -> bool: ...
151-
def __lt__(self, other: timedelta) -> bool: ...
152-
def __ge__(self, other: timedelta) -> bool: ...
153-
def __gt__(self, other: timedelta) -> bool: ...
154-
def __hash__(self) -> int: ...
353+
# Mypy complains Forward operator "<inequality op>" is not callable, so ignore misc
354+
# for le, lt ge and gt
355+
# Override due to more types supported than dt.timedelta
356+
@overload # type: ignore[override]
357+
def __le__(self, other: timedelta | Timedelta | np.timedelta64) -> bool: ... # type: ignore[misc]
358+
@overload
359+
def __le__(
360+
self, other: TimedeltaIndex | npt.NDArray[np.timedelta64]
361+
) -> npt.NDArray[np.bool_]: ...
362+
@overload
363+
def __le__(self, other: TimedeltaSeries | Series[pd.Timedelta]) -> Series[bool]: ...
364+
# Override due to more types supported than dt.timedelta
365+
@overload # type: ignore[override]
366+
def __lt__(self, other: timedelta | Timedelta | np.timedelta64) -> bool: ... # type: ignore[misc]
367+
@overload
368+
def __lt__(
369+
self, other: TimedeltaIndex | npt.NDArray[np.timedelta64]
370+
) -> npt.NDArray[np.bool_]: ...
371+
@overload
372+
def __lt__(self, other: TimedeltaSeries | Series[pd.Timedelta]) -> Series[bool]: ...
373+
# Override due to more types supported than dt.timedelta
374+
@overload # type: ignore[override]
375+
def __ge__(self, other: timedelta | Timedelta | np.timedelta64) -> bool: ... # type: ignore[misc]
376+
@overload
377+
def __ge__(
378+
self, other: TimedeltaIndex | npt.NDArray[np.timedelta64]
379+
) -> npt.NDArray[np.bool_]: ...
380+
@overload
381+
def __ge__(self, other: TimedeltaSeries | Series[pd.Timedelta]) -> Series[bool]: ...
382+
# Override due to more types supported than dt.timedelta
383+
@overload # type: ignore[override]
384+
def __gt__(self, other: timedelta | Timedelta | np.timedelta64) -> bool: ... # type: ignore[misc]
385+
@overload
386+
def __gt__(
387+
self, other: TimedeltaIndex | npt.NDArray[np.timedelta64]
388+
) -> npt.NDArray[np.bool_]: ...
389+
@overload
390+
def __gt__(self, other: TimedeltaSeries | Series[pd.Timedelta]) -> Series[bool]: ...
155391
def isoformat(self) -> str: ...
156392
def to_numpy(self) -> np.timedelta64: ...
393+
@property
394+
def components(self) -> Components: ...
395+
def view(self, dtype: npt.DTypeLike = ...) -> object: ...

0 commit comments

Comments
 (0)