Skip to content

Commit 0067f7c

Browse files
cmp0xffcmp0xff
committed
feat(arithmetic): multiplication
1 parent b197a8b commit 0067f7c

File tree

22 files changed

+728
-520
lines changed

22 files changed

+728
-520
lines changed

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

Lines changed: 53 additions & 86 deletions
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,9 @@
11
# pyright: strict
2-
import datetime as dt
3-
from datetime import timedelta
2+
from datetime import (
3+
date,
4+
datetime,
5+
timedelta,
6+
)
47
from typing import (
58
ClassVar,
69
Literal,
@@ -33,6 +36,11 @@ from pandas._typing import (
3336
TimeUnit,
3437
np_1darray,
3538
np_ndarray,
39+
np_ndarray_anyint,
40+
np_ndarray_bool,
41+
np_ndarray_dt,
42+
np_ndarray_float,
43+
np_ndarray_td,
3644
npt,
3745
)
3846

@@ -135,58 +143,48 @@ class Timedelta(timedelta):
135143
def ceil(self, freq: str | BaseOffset) -> Self: ...
136144
@property
137145
def resolution_string(self) -> str: ...
138-
# Override due to more types supported than dt.timedelta
146+
# Override due to more types supported than timedelta
139147
@overload # type: ignore[override]
140-
def __add__(self, other: dt.datetime | np.datetime64) -> Timestamp: ...
148+
def __add__(self, other: datetime | np.datetime64) -> Timestamp: ...
141149
@overload
142150
def __add__(self, other: timedelta | np.timedelta64) -> Self: ...
143151
@overload
144152
def __add__(self, other: NaTType) -> NaTType: ...
145153
@overload
146154
def __add__(self, other: Period) -> Period: ...
147155
@overload
148-
def __add__(self, other: dt.date) -> dt.date: ...
156+
def __add__(self, other: date) -> date: ...
149157
@overload
150-
def __add__(
151-
self, other: np_ndarray[ShapeT, np.timedelta64]
152-
) -> np_ndarray[ShapeT, np.timedelta64]: ...
158+
def __add__(self, other: np_ndarray_td) -> np_ndarray_td: ...
153159
@overload
154-
def __add__(
155-
self, other: np_ndarray[ShapeT, np.datetime64]
156-
) -> np_ndarray[ShapeT, np.datetime64]: ...
160+
def __add__(self, other: np_ndarray_dt) -> np_ndarray_dt: ...
157161
@overload
158-
def __radd__(self, other: dt.datetime | np.datetime64) -> Timestamp: ... # type: ignore[misc]
162+
def __radd__(self, other: datetime | np.datetime64) -> Timestamp: ... # type: ignore[misc]
159163
@overload
160164
def __radd__(self, other: timedelta | np.timedelta64) -> Self: ...
161165
@overload
162166
def __radd__(self, other: NaTType) -> NaTType: ...
163167
@overload
164-
def __radd__(self, other: dt.date) -> dt.date: ...
168+
def __radd__(self, other: date) -> date: ...
165169
@overload
166-
def __radd__(
167-
self, other: np_ndarray[ShapeT, np.timedelta64]
168-
) -> np_ndarray[ShapeT, np.timedelta64]: ...
170+
def __radd__(self, other: np_ndarray_td) -> np_ndarray_td: ...
169171
@overload
170-
def __radd__(
171-
self, other: np_ndarray[ShapeT, np.datetime64]
172-
) -> np_ndarray[ShapeT, np.datetime64]: ...
173-
# Override due to more types supported than dt.timedelta
172+
def __radd__(self, other: np_ndarray_dt) -> np_ndarray_dt: ...
173+
# Override due to more types supported than timedelta
174174
@overload # type: ignore[override]
175175
def __sub__(self, other: timedelta | Timedelta | np.timedelta64) -> Timedelta: ...
176176
@overload
177177
def __sub__(self, other: NaTType) -> NaTType: ...
178178
@overload
179-
def __sub__(
180-
self, other: np_ndarray[ShapeT, np.timedelta64]
181-
) -> np_ndarray[ShapeT, np.timedelta64]: ...
179+
def __sub__(self, other: np_ndarray_td) -> np_ndarray_td: ...
182180
@overload
183181
def __sub__(self, other: pd.TimedeltaIndex) -> TimedeltaIndex: ...
184182
@overload
185183
def __sub__(self, other: Series[pd.Timedelta]) -> Series[pd.Timedelta]: ...
186184
@overload
187185
def __rsub__(self, other: timedelta | Timedelta | np.timedelta64) -> Timedelta: ...
188186
@overload
189-
def __rsub__(self, other: dt.datetime | Timestamp | np.datetime64) -> Timestamp: ... # type: ignore[misc]
187+
def __rsub__(self, other: datetime | Timestamp | np.datetime64) -> Timestamp: ... # type: ignore[misc]
190188
@overload
191189
def __rsub__(self, other: NaTType) -> NaTType: ...
192190
@overload
@@ -196,58 +194,39 @@ class Timedelta(timedelta):
196194
@overload
197195
def __rsub__(self, other: DatetimeIndex) -> DatetimeIndex: ...
198196
@overload
199-
def __rsub__(
200-
self, other: np_ndarray[ShapeT, np.datetime64]
201-
) -> np_ndarray[ShapeT, np.datetime64]: ...
197+
def __rsub__(self, other: np_ndarray_dt) -> np_ndarray_dt: ...
202198
@overload
203-
def __rsub__(
204-
self, other: np_ndarray[ShapeT, np.timedelta64]
205-
) -> np_ndarray[ShapeT, np.timedelta64]: ...
199+
def __rsub__(self, other: np_ndarray_td) -> np_ndarray_td: ...
206200
@overload
207201
def __rsub__(self, other: pd.TimedeltaIndex) -> pd.TimedeltaIndex: ...
208202
def __neg__(self) -> Timedelta: ...
209203
def __pos__(self) -> Timedelta: ...
210204
def __abs__(self) -> Timedelta: ...
211-
# Override due to more types supported than dt.timedelta
205+
# Override due to more types supported than timedelta
212206
@overload # type: ignore[override]
213207
def __mul__(self, other: float) -> Timedelta: ...
214208
@overload
215209
def __mul__(
216-
self, other: np_ndarray[ShapeT, np.integer] | np_ndarray[ShapeT, np.floating]
217-
) -> np_ndarray[ShapeT, np.timedelta64]: ...
218-
@overload
219-
def __mul__(self, other: Series[int]) -> Series[Timedelta]: ...
220-
@overload
221-
def __mul__(self, other: Series[float]) -> Series[Timedelta]: ...
222-
@overload
223-
def __mul__(self, other: Index[int] | Index[float]) -> TimedeltaIndex: ...
210+
self, other: np_ndarray_bool | np_ndarray_anyint | np_ndarray_float
211+
) -> np_ndarray_td: ...
224212
@overload
225213
def __rmul__(self, other: float) -> Timedelta: ...
226214
@overload
227215
def __rmul__(
228-
self, other: np_ndarray[ShapeT, np.floating] | np_ndarray[ShapeT, np.integer]
229-
) -> np_ndarray[ShapeT, np.timedelta64]: ...
230-
@overload
231-
def __rmul__(self, other: Series[int]) -> Series[Timedelta]: ...
232-
@overload
233-
def __rmul__(self, other: Series[float]) -> Series[Timedelta]: ...
234-
# maybe related to https://github.com/python/mypy/issues/10755
235-
@overload
236-
def __rmul__(self, other: Index[int] | Index[float]) -> TimedeltaIndex: ...
237-
# Override due to more types supported than dt.timedelta
216+
self, other: np_ndarray_bool | np_ndarray_anyint | np_ndarray_float
217+
) -> np_ndarray_td: ...
218+
# Override due to more types supported than timedelta
238219
# error: Signature of "__floordiv__" incompatible with supertype "timedelta"
239220
@overload # type: ignore[override]
240221
def __floordiv__(self, other: timedelta | Timedelta | np.timedelta64) -> int: ...
241222
@overload
242223
def __floordiv__(self, other: float) -> Timedelta: ...
243224
@overload
244225
def __floordiv__(
245-
self, other: np_ndarray[ShapeT, np.integer] | np_ndarray[ShapeT, np.floating]
246-
) -> np_ndarray[ShapeT, np.timedelta64]: ...
226+
self, other: np_ndarray_anyint | np_ndarray_float
227+
) -> np_ndarray_td: ...
247228
@overload
248-
def __floordiv__(
249-
self, other: np_ndarray[ShapeT, np.timedelta64]
250-
) -> np_ndarray[ShapeT, np.int_]: ...
229+
def __floordiv__(self, other: np_ndarray_td) -> np_ndarray[ShapeT, np.int_]: ...
251230
@overload
252231
def __floordiv__(self, other: Index[int] | Index[float]) -> TimedeltaIndex: ...
253232
@overload
@@ -263,18 +242,16 @@ class Timedelta(timedelta):
263242
@overload
264243
def __rfloordiv__(self, other: NaTType | None) -> float: ...
265244
@overload
266-
def __rfloordiv__(
267-
self, other: np_ndarray[ShapeT, np.timedelta64]
268-
) -> np_ndarray[ShapeT, np.int_]: ...
269-
# Override due to more types supported than dt.timedelta
245+
def __rfloordiv__(self, other: np_ndarray_td) -> np_ndarray[ShapeT, np.int_]: ...
246+
# Override due to more types supported than timedelta
270247
@overload # type: ignore[override]
271248
def __truediv__(self, other: timedelta | Timedelta | NaTType) -> float: ...
272249
@overload
273250
def __truediv__(self, other: float) -> Timedelta: ...
274251
@overload
275252
def __truediv__(
276-
self, other: np_ndarray[ShapeT, np.integer] | np_ndarray[ShapeT, np.floating]
277-
) -> np_ndarray[ShapeT, np.timedelta64]: ...
253+
self, other: np_ndarray_anyint | np_ndarray_float
254+
) -> np_ndarray_td: ...
278255
@overload
279256
def __truediv__(self, other: Series[Timedelta]) -> Series[float]: ...
280257
@overload
@@ -284,7 +261,7 @@ class Timedelta(timedelta):
284261
@overload
285262
def __truediv__(self, other: Index[int] | Index[float]) -> TimedeltaIndex: ...
286263
def __rtruediv__(self, other: timedelta | Timedelta | NaTType) -> float: ...
287-
# Override due to more types supported than dt.timedelta
264+
# Override due to more types supported than timedelta
288265
@overload
289266
def __eq__(self, other: timedelta | Timedelta | np.timedelta64) -> bool: ... # type: ignore[overload-overlap] # pyright: ignore[reportOverlappingOverload]
290267
@overload
@@ -293,11 +270,11 @@ class Timedelta(timedelta):
293270
def __eq__(self, other: Index) -> np_1darray[np.bool]: ... # type: ignore[overload-overlap]
294271
@overload
295272
def __eq__( # type: ignore[overload-overlap]
296-
self, other: np_ndarray[ShapeT, np.timedelta64]
297-
) -> np_ndarray[ShapeT, np.bool_]: ...
273+
self, other: np_ndarray_td
274+
) -> np_ndarray_bool: ...
298275
@overload
299276
def __eq__(self, other: object) -> Literal[False]: ...
300-
# Override due to more types supported than dt.timedelta
277+
# Override due to more types supported than timedelta
301278
@overload
302279
def __ne__(self, other: timedelta | Timedelta | np.timedelta64) -> bool: ... # type: ignore[overload-overlap] # pyright: ignore[reportOverlappingOverload]
303280
@overload
@@ -306,11 +283,11 @@ class Timedelta(timedelta):
306283
def __ne__(self, other: Index) -> np_1darray[np.bool]: ... # type: ignore[overload-overlap]
307284
@overload
308285
def __ne__( # type: ignore[overload-overlap]
309-
self, other: np_ndarray[ShapeT, np.timedelta64]
310-
) -> np_ndarray[ShapeT, np.bool_]: ...
286+
self, other: np_ndarray_td
287+
) -> np_ndarray_bool: ...
311288
@overload
312289
def __ne__(self, other: object) -> Literal[True]: ...
313-
# Override due to more types supported than dt.timedelta
290+
# Override due to more types supported than timedelta
314291
@overload # type: ignore[override]
315292
def __mod__(self, other: timedelta) -> Timedelta: ...
316293
@overload
@@ -320,58 +297,48 @@ class Timedelta(timedelta):
320297
@overload
321298
def __mod__(self, other: Index[int] | Index[float]) -> TimedeltaIndex: ...
322299
@overload
323-
def __mod__(
324-
self, other: np_ndarray[ShapeT, np.integer] | np_ndarray[ShapeT, np.floating]
325-
) -> np_ndarray[ShapeT, np.timedelta64]: ...
300+
def __mod__(self, other: np_ndarray_anyint | np_ndarray_float) -> np_ndarray_td: ...
326301
@overload
327302
def __mod__(
328303
self, other: Series[int] | Series[float] | Series[Timedelta]
329304
) -> Series[Timedelta]: ...
330305
def __divmod__(self, other: timedelta) -> tuple[int, Timedelta]: ...
331306
# Mypy complains Forward operator "<inequality op>" is not callable, so ignore misc
332307
# for le, lt ge and gt
333-
# Override due to more types supported than dt.timedelta
308+
# Override due to more types supported than timedelta
334309
@overload # type: ignore[override]
335310
def __le__(self, other: timedelta | Timedelta | np.timedelta64) -> bool: ... # type: ignore[misc]
336311
@overload
337312
def __le__(self, other: TimedeltaIndex) -> np_1darray[np.bool]: ...
338313
@overload
339-
def __le__(
340-
self, other: np_ndarray[ShapeT, np.timedelta64]
341-
) -> np_ndarray[ShapeT, np.bool_]: ...
314+
def __le__(self, other: np_ndarray_td) -> np_ndarray_bool: ...
342315
@overload
343316
def __le__(self, other: Series[pd.Timedelta]) -> Series[bool]: ...
344-
# Override due to more types supported than dt.timedelta
317+
# Override due to more types supported than timedelta
345318
@overload # type: ignore[override]
346319
def __lt__(self, other: timedelta | Timedelta | np.timedelta64) -> bool: ... # type: ignore[misc]
347320
@overload
348321
def __lt__(self, other: TimedeltaIndex) -> np_1darray[np.bool]: ...
349322
@overload
350-
def __lt__(
351-
self, other: np_ndarray[ShapeT, np.timedelta64]
352-
) -> np_ndarray[ShapeT, np.bool_]: ...
323+
def __lt__(self, other: np_ndarray_td) -> np_ndarray_bool: ...
353324
@overload
354325
def __lt__(self, other: Series[pd.Timedelta]) -> Series[bool]: ...
355-
# Override due to more types supported than dt.timedelta
326+
# Override due to more types supported than timedelta
356327
@overload # type: ignore[override]
357328
def __ge__(self, other: timedelta | Timedelta | np.timedelta64) -> bool: ... # type: ignore[misc]
358329
@overload
359330
def __ge__(self, other: TimedeltaIndex) -> np_1darray[np.bool]: ...
360331
@overload
361-
def __ge__(
362-
self, other: np_ndarray[ShapeT, np.timedelta64]
363-
) -> np_ndarray[ShapeT, np.bool_]: ...
332+
def __ge__(self, other: np_ndarray_td) -> np_ndarray_bool: ...
364333
@overload
365334
def __ge__(self, other: Series[pd.Timedelta]) -> Series[bool]: ...
366-
# Override due to more types supported than dt.timedelta
335+
# Override due to more types supported than timedelta
367336
@overload # type: ignore[override]
368337
def __gt__(self, other: timedelta | Timedelta | np.timedelta64) -> bool: ... # type: ignore[misc]
369338
@overload
370339
def __gt__(self, other: TimedeltaIndex) -> np_1darray[np.bool]: ...
371340
@overload
372-
def __gt__(
373-
self, other: np_ndarray[ShapeT, np.timedelta64]
374-
) -> np_ndarray[ShapeT, np.bool_]: ...
341+
def __gt__(self, other: np_ndarray_td) -> np_ndarray_bool: ...
375342
@overload
376343
def __gt__(self, other: Series[pd.Timedelta]) -> Series[bool]: ...
377344
def __hash__(self) -> int: ...

‎pandas-stubs/_typing.pyi‎

Lines changed: 12 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -173,7 +173,8 @@ RandomState: TypeAlias = (
173173
)
174174

175175
# dtypes
176-
NpDtype: TypeAlias = str | np.dtype[np.generic] | type[str | complex | bool | object]
176+
NpDtypeNoStr: TypeAlias = np.dtype[np.generic] | type[complex | bool | object]
177+
NpDtype: TypeAlias = str | NpDtypeNoStr | type[str]
177178
Dtype: TypeAlias = ExtensionDtype | NpDtype
178179

179180
# AstypeArg is more carefully defined here as compared to pandas
@@ -847,19 +848,21 @@ MaskType: TypeAlias = Series[bool] | np_ndarray_bool | list[bool]
847848

848849
T_INT = TypeVar("T_INT", bound=int)
849850
T_COMPLEX = TypeVar("T_COMPLEX", bound=complex)
850-
SeriesDTypeNoDateTime: TypeAlias = (
851-
str
852-
| bytes
851+
SeriesDTypeNoStrDateTime: TypeAlias = (
852+
bytes
853853
| bool
854854
| int
855855
| float
856856
| complex
857-
| Dtype
857+
| NpDtypeNoStr
858+
| ExtensionDtype
858859
| Period
859860
| Interval
860861
| CategoricalDtype
861862
| BaseOffset
862-
| list[str]
863+
)
864+
SeriesDTypeNoDateTime: TypeAlias = (
865+
str | SeriesDTypeNoStrDateTime | type[str] | list[str]
863866
)
864867
SeriesDType: TypeAlias = (
865868
SeriesDTypeNoDateTime
@@ -869,6 +872,9 @@ SeriesDType: TypeAlias = (
869872
| datetime.timedelta # includes pd.Timedelta
870873
)
871874
S1 = TypeVar("S1", bound=SeriesDType, default=Any)
875+
S1_CO_NSDT = TypeVar(
876+
"S1_CO_NSDT", bound=SeriesDTypeNoStrDateTime, default=Any, covariant=True
877+
)
872878
S1_CT_NDT = TypeVar(
873879
"S1_CT_NDT", bound=SeriesDTypeNoDateTime, default=Any, contravariant=True
874880
)

0 commit comments

Comments
 (0)