Skip to content

Commit 519be94

Browse files
hamdanalhamdanal
committed
Type timedelta array shapes and refactor ts and td tests
1 parent abfef56 commit 519be94

File tree

2 files changed

+432
-518
lines changed

2 files changed

+432
-518
lines changed

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

Lines changed: 55 additions & 40 deletions
Original file line numberDiff line numberDiff line change
@@ -33,7 +33,10 @@ from pandas._libs.tslibs import (
3333
from pandas._libs.tslibs.period import Period
3434
from pandas._libs.tslibs.timestamps import Timestamp
3535
from pandas._typing import (
36+
ShapeT,
3637
TimeUnit,
38+
np_1darray,
39+
np_ndarray,
3740
npt,
3841
)
3942

@@ -153,12 +156,12 @@ class Timedelta(timedelta):
153156
def __add__(self, other: DatetimeIndex) -> DatetimeIndex: ...
154157
@overload
155158
def __add__(
156-
self, other: npt.NDArray[np.timedelta64]
157-
) -> npt.NDArray[np.timedelta64]: ...
159+
self, other: np_ndarray[ShapeT, np.timedelta64]
160+
) -> np_ndarray[ShapeT, np.timedelta64]: ...
158161
@overload
159162
def __add__(
160-
self, other: npt.NDArray[np.datetime64]
161-
) -> npt.NDArray[np.datetime64]: ...
163+
self, other: np_ndarray[ShapeT, np.datetime64]
164+
) -> np_ndarray[ShapeT, np.datetime64]: ...
162165
@overload
163166
def __add__(self, other: pd.TimedeltaIndex) -> pd.TimedeltaIndex: ...
164167
@overload
@@ -176,12 +179,12 @@ class Timedelta(timedelta):
176179
def __radd__(self, other: NaTType) -> NaTType: ...
177180
@overload
178181
def __radd__(
179-
self, other: npt.NDArray[np.timedelta64]
180-
) -> npt.NDArray[np.timedelta64]: ...
182+
self, other: np_ndarray[ShapeT, np.timedelta64]
183+
) -> np_ndarray[ShapeT, np.timedelta64]: ...
181184
@overload
182185
def __radd__(
183-
self, other: npt.NDArray[np.datetime64]
184-
) -> npt.NDArray[np.datetime64]: ...
186+
self, other: np_ndarray[ShapeT, np.datetime64]
187+
) -> np_ndarray[ShapeT, np.datetime64]: ...
185188
@overload
186189
def __radd__(self, other: pd.TimedeltaIndex) -> pd.TimedeltaIndex: ...
187190
@overload
@@ -193,8 +196,8 @@ class Timedelta(timedelta):
193196
def __sub__(self, other: NaTType) -> NaTType: ...
194197
@overload
195198
def __sub__(
196-
self, other: npt.NDArray[np.timedelta64]
197-
) -> npt.NDArray[np.timedelta64]: ...
199+
self, other: np_ndarray[ShapeT, np.timedelta64]
200+
) -> np_ndarray[ShapeT, np.timedelta64]: ...
198201
@overload
199202
def __sub__(self, other: pd.TimedeltaIndex) -> TimedeltaIndex: ...
200203
@overload
@@ -215,12 +218,12 @@ class Timedelta(timedelta):
215218
def __rsub__(self, other: DatetimeIndex) -> DatetimeIndex: ...
216219
@overload
217220
def __rsub__(
218-
self, other: npt.NDArray[np.datetime64]
219-
) -> npt.NDArray[np.datetime64]: ...
221+
self, other: np_ndarray[ShapeT, np.datetime64]
222+
) -> np_ndarray[ShapeT, np.datetime64]: ...
220223
@overload
221224
def __rsub__(
222-
self, other: npt.NDArray[np.timedelta64]
223-
) -> npt.NDArray[np.timedelta64]: ...
225+
self, other: np_ndarray[ShapeT, np.timedelta64]
226+
) -> np_ndarray[ShapeT, np.timedelta64]: ...
224227
@overload
225228
def __rsub__(self, other: pd.TimedeltaIndex) -> pd.TimedeltaIndex: ...
226229
def __neg__(self) -> Timedelta: ...
@@ -231,8 +234,8 @@ class Timedelta(timedelta):
231234
def __mul__(self, other: float) -> Timedelta: ...
232235
@overload
233236
def __mul__(
234-
self, other: npt.NDArray[np.integer] | npt.NDArray[np.floating]
235-
) -> npt.NDArray[np.timedelta64]: ...
237+
self, other: np_ndarray[ShapeT, np.integer] | np_ndarray[ShapeT, np.floating]
238+
) -> np_ndarray[ShapeT, np.timedelta64]: ...
236239
@overload
237240
def __mul__(self, other: Series[int]) -> TimedeltaSeries: ...
238241
@overload
@@ -243,8 +246,8 @@ class Timedelta(timedelta):
243246
def __rmul__(self, other: float) -> Timedelta: ...
244247
@overload
245248
def __rmul__(
246-
self, other: npt.NDArray[np.floating] | npt.NDArray[np.integer]
247-
) -> npt.NDArray[np.timedelta64]: ...
249+
self, other: np_ndarray[ShapeT, np.floating] | np_ndarray[ShapeT, np.integer]
250+
) -> np_ndarray[ShapeT, np.timedelta64]: ...
248251
@overload
249252
def __rmul__(self, other: Series[int]) -> TimedeltaSeries: ...
250253
@overload
@@ -260,12 +263,12 @@ class Timedelta(timedelta):
260263
def __floordiv__(self, other: float) -> Timedelta: ...
261264
@overload
262265
def __floordiv__(
263-
self, other: npt.NDArray[np.integer] | npt.NDArray[np.floating]
264-
) -> npt.NDArray[np.timedelta64]: ...
266+
self, other: np_ndarray[ShapeT, np.integer] | np_ndarray[ShapeT, np.floating]
267+
) -> np_ndarray[ShapeT, np.timedelta64]: ...
265268
@overload
266269
def __floordiv__(
267-
self, other: npt.NDArray[np.timedelta64]
268-
) -> npt.NDArray[np.int_]: ...
270+
self, other: np_ndarray[ShapeT, np.timedelta64]
271+
) -> np_ndarray[ShapeT, np.int_]: ...
269272
@overload
270273
def __floordiv__(self, other: Index[int] | Index[float]) -> TimedeltaIndex: ...
271274
@overload
@@ -282,17 +285,17 @@ class Timedelta(timedelta):
282285
def __rfloordiv__(self, other: NaTType | None) -> float: ...
283286
@overload
284287
def __rfloordiv__(
285-
self, other: npt.NDArray[np.timedelta64]
286-
) -> npt.NDArray[np.int_]: ...
288+
self, other: np_ndarray[ShapeT, np.timedelta64]
289+
) -> np_ndarray[ShapeT, np.int_]: ...
287290
# Override due to more types supported than dt.timedelta
288291
@overload # type: ignore[override]
289292
def __truediv__(self, other: timedelta | Timedelta | NaTType) -> float: ...
290293
@overload
291294
def __truediv__(self, other: float) -> Timedelta: ...
292295
@overload
293296
def __truediv__(
294-
self, other: npt.NDArray[np.integer] | npt.NDArray[np.floating]
295-
) -> npt.NDArray[np.timedelta64]: ...
297+
self, other: np_ndarray[ShapeT, np.integer] | np_ndarray[ShapeT, np.floating]
298+
) -> np_ndarray[ShapeT, np.timedelta64]: ...
296299
@overload
297300
def __truediv__(self, other: TimedeltaSeries) -> Series[float]: ...
298301
@overload
@@ -308,9 +311,11 @@ class Timedelta(timedelta):
308311
@overload
309312
def __eq__(self, other: TimedeltaSeries | Series[pd.Timedelta]) -> Series[bool]: ... # type: ignore[overload-overlap]
310313
@overload
314+
def __eq__(self, other: TimedeltaIndex) -> np_1darray[np.bool]: ... # type: ignore[overload-overlap]
315+
@overload
311316
def __eq__( # type: ignore[overload-overlap]
312-
self, other: TimedeltaIndex | npt.NDArray[np.timedelta64]
313-
) -> npt.NDArray[np.bool_]: ...
317+
self, other: np_ndarray[ShapeT, np.timedelta64]
318+
) -> np_ndarray[ShapeT, np.bool_]: ...
314319
@overload
315320
def __eq__(self, other: object) -> Literal[False]: ...
316321
# Override due to more types supported than dt.timedelta
@@ -319,9 +324,11 @@ class Timedelta(timedelta):
319324
@overload
320325
def __ne__(self, other: TimedeltaSeries | Series[pd.Timedelta]) -> Series[bool]: ... # type: ignore[overload-overlap]
321326
@overload
327+
def __ne__(self, other: TimedeltaIndex) -> np_1darray[np.bool]: ... # type: ignore[overload-overlap]
328+
@overload
322329
def __ne__( # type: ignore[overload-overlap]
323-
self, other: TimedeltaIndex | npt.NDArray[np.timedelta64]
324-
) -> npt.NDArray[np.bool_]: ...
330+
self, other: np_ndarray[ShapeT, np.timedelta64]
331+
) -> np_ndarray[ShapeT, np.bool_]: ...
325332
@overload
326333
def __ne__(self, other: object) -> Literal[True]: ...
327334
# Override due to more types supported than dt.timedelta
@@ -335,8 +342,8 @@ class Timedelta(timedelta):
335342
def __mod__(self, other: Index[int] | Index[float]) -> TimedeltaIndex: ...
336343
@overload
337344
def __mod__(
338-
self, other: npt.NDArray[np.integer] | npt.NDArray[np.floating]
339-
) -> npt.NDArray[np.timedelta64]: ...
345+
self, other: np_ndarray[ShapeT, np.integer] | np_ndarray[ShapeT, np.floating]
346+
) -> np_ndarray[ShapeT, np.timedelta64]: ...
340347
@overload
341348
def __mod__(
342349
self, other: Series[int] | Series[float] | TimedeltaSeries
@@ -348,36 +355,44 @@ class Timedelta(timedelta):
348355
@overload # type: ignore[override]
349356
def __le__(self, other: timedelta | Timedelta | np.timedelta64) -> bool: ... # type: ignore[misc]
350357
@overload
358+
def __le__(self, other: TimedeltaIndex) -> np_1darray[np.bool]: ...
359+
@overload
351360
def __le__(
352-
self, other: TimedeltaIndex | npt.NDArray[np.timedelta64]
353-
) -> npt.NDArray[np.bool_]: ...
361+
self, other: np_ndarray[ShapeT, np.timedelta64]
362+
) -> np_ndarray[ShapeT, np.bool_]: ...
354363
@overload
355364
def __le__(self, other: TimedeltaSeries | Series[pd.Timedelta]) -> Series[bool]: ...
356365
# Override due to more types supported than dt.timedelta
357366
@overload # type: ignore[override]
358367
def __lt__(self, other: timedelta | Timedelta | np.timedelta64) -> bool: ... # type: ignore[misc]
359368
@overload
369+
def __lt__(self, other: TimedeltaIndex) -> np_1darray[np.bool]: ...
370+
@overload
360371
def __lt__(
361-
self, other: TimedeltaIndex | npt.NDArray[np.timedelta64]
362-
) -> npt.NDArray[np.bool_]: ...
372+
self, other: np_ndarray[ShapeT, np.timedelta64]
373+
) -> np_ndarray[ShapeT, np.bool_]: ...
363374
@overload
364375
def __lt__(self, other: TimedeltaSeries | Series[pd.Timedelta]) -> Series[bool]: ...
365376
# Override due to more types supported than dt.timedelta
366377
@overload # type: ignore[override]
367378
def __ge__(self, other: timedelta | Timedelta | np.timedelta64) -> bool: ... # type: ignore[misc]
368379
@overload
380+
def __ge__(self, other: TimedeltaIndex) -> np_1darray[np.bool]: ...
381+
@overload
369382
def __ge__(
370-
self, other: TimedeltaIndex | npt.NDArray[np.timedelta64]
371-
) -> npt.NDArray[np.bool_]: ...
383+
self, other: np_ndarray[ShapeT, np.timedelta64]
384+
) -> np_ndarray[ShapeT, np.bool_]: ...
372385
@overload
373386
def __ge__(self, other: TimedeltaSeries | Series[pd.Timedelta]) -> Series[bool]: ...
374387
# Override due to more types supported than dt.timedelta
375388
@overload # type: ignore[override]
376389
def __gt__(self, other: timedelta | Timedelta | np.timedelta64) -> bool: ... # type: ignore[misc]
377390
@overload
391+
def __gt__(self, other: TimedeltaIndex) -> np_1darray[np.bool]: ...
392+
@overload
378393
def __gt__(
379-
self, other: TimedeltaIndex | npt.NDArray[np.timedelta64]
380-
) -> npt.NDArray[np.bool_]: ...
394+
self, other: np_ndarray[ShapeT, np.timedelta64]
395+
) -> np_ndarray[ShapeT, np.bool_]: ...
381396
@overload
382397
def __gt__(self, other: TimedeltaSeries | Series[pd.Timedelta]) -> Series[bool]: ...
383398
def __hash__(self) -> int: ...

0 commit comments

Comments
 (0)