Skip to content

Commit 860b2bf

Browse files
committed
feat(arithmetic): multiplication
1 parent 0042dfd commit 860b2bf

File tree

23 files changed

+860
-507
lines changed

23 files changed

+860
-507
lines changed

pandas-stubs/_libs/tslibs/period.pyi

Lines changed: 16 additions & 17 deletions
Original file line numberDiff line numberDiff line change
@@ -12,7 +12,10 @@ from pandas import (
1212
Timedelta,
1313
TimedeltaIndex,
1414
)
15-
from typing_extensions import TypeAlias
15+
from typing_extensions import (
16+
Self,
17+
TypeAlias,
18+
)
1619

1720
from pandas._libs.tslibs import NaTType
1821
from pandas._libs.tslibs.offsets import BaseOffset
@@ -87,15 +90,23 @@ class Period(PeriodMixin):
8790
@overload
8891
def __sub__(self, other: TimedeltaIndex) -> PeriodIndex: ...
8992
@overload
90-
def __add__(self, other: _PeriodAddSub) -> Period: ...
93+
def __add__(self, other: _PeriodAddSub) -> Self: ...
9194
@overload
9295
def __add__(self, other: NaTType) -> NaTType: ...
9396
@overload
9497
def __add__(self, other: Index) -> PeriodIndex: ...
98+
# Ignored due to indecipherable error from mypy:
99+
# Forward operator "__add__" is not callable [misc]
95100
@overload
96-
def __add__(
97-
self, other: Series[BaseOffset] | Series[Timedelta]
98-
) -> Series[Period]: ... # pyrefly: ignore[bad-specialization]
101+
def __radd__(self, other: _PeriodAddSub) -> Self: ... # type: ignore[misc]
102+
@overload
103+
def __radd__(self, other: NaTType) -> NaTType: ...
104+
# Real signature is -> PeriodIndex, but conflicts with Index.__add__
105+
# Changing Index is very hard due to Index inheritance
106+
# Signatures of "__radd__" of "Period" and "__add__" of "Index"
107+
# are unsafely overlapping
108+
@overload
109+
def __radd__(self, other: Index) -> PeriodIndex: ...
99110
# ignore[misc] here because we know all other comparisons
100111
# are False, so we use Literal[False]
101112
@overload
@@ -168,18 +179,6 @@ class Period(PeriodMixin):
168179
def __ne__(self, other: np_ndarray[ShapeT, np.object_]) -> np_ndarray[ShapeT, np.bool]: ... # type: ignore[overload-overlap]
169180
@overload
170181
def __ne__(self, other: object) -> Literal[True]: ...
171-
# Ignored due to indecipherable error from mypy:
172-
# Forward operator "__add__" is not callable [misc]
173-
@overload
174-
def __radd__(self, other: _PeriodAddSub) -> Period: ... # type: ignore[misc]
175-
# Real signature is -> PeriodIndex, but conflicts with Index.__add__
176-
# Changing Index is very hard due to Index inheritance
177-
# Signatures of "__radd__" of "Period" and "__add__" of "Index"
178-
# are unsafely overlapping
179-
@overload
180-
def __radd__(self, other: Index) -> Index: ...
181-
@overload
182-
def __radd__(self, other: NaTType) -> NaTType: ...
183182
@property
184183
def day(self) -> int: ...
185184
@property

pandas-stubs/_libs/tslibs/timedeltas.pyi

Lines changed: 38 additions & 50 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,
@@ -135,17 +138,17 @@ class Timedelta(timedelta):
135138
def ceil(self, freq: str | BaseOffset) -> Self: ...
136139
@property
137140
def resolution_string(self) -> str: ...
138-
# Override due to more types supported than dt.timedelta
141+
# Override due to more types supported than timedelta
139142
@overload # type: ignore[override]
140-
def __add__(self, other: dt.datetime | np.datetime64) -> Timestamp: ...
143+
def __add__(self, other: datetime | np.datetime64) -> Timestamp: ...
141144
@overload
142145
def __add__(self, other: timedelta | np.timedelta64) -> Self: ...
143146
@overload
144147
def __add__(self, other: NaTType) -> NaTType: ...
145148
@overload
146149
def __add__(self, other: Period) -> Period: ...
147150
@overload
148-
def __add__(self, other: dt.date) -> dt.date: ...
151+
def __add__(self, other: date) -> date: ...
149152
@overload
150153
def __add__(
151154
self, other: np_ndarray[ShapeT, np.timedelta64]
@@ -155,13 +158,13 @@ class Timedelta(timedelta):
155158
self, other: np_ndarray[ShapeT, np.datetime64]
156159
) -> np_ndarray[ShapeT, np.datetime64]: ...
157160
@overload
158-
def __radd__(self, other: dt.datetime | np.datetime64) -> Timestamp: ... # type: ignore[misc]
161+
def __radd__(self, other: datetime | np.datetime64) -> Timestamp: ... # type: ignore[misc]
159162
@overload
160163
def __radd__(self, other: timedelta | np.timedelta64) -> Self: ...
161164
@overload
162165
def __radd__(self, other: NaTType) -> NaTType: ...
163166
@overload
164-
def __radd__(self, other: dt.date) -> dt.date: ...
167+
def __radd__(self, other: date) -> date: ...
165168
@overload
166169
def __radd__(
167170
self, other: np_ndarray[ShapeT, np.timedelta64]
@@ -170,9 +173,9 @@ class Timedelta(timedelta):
170173
def __radd__(
171174
self, other: np_ndarray[ShapeT, np.datetime64]
172175
) -> np_ndarray[ShapeT, np.datetime64]: ...
173-
# Override due to more types supported than dt.timedelta
176+
# Override due to more types supported than timedelta
174177
@overload # type: ignore[override]
175-
def __sub__(self, other: timedelta | Timedelta | np.timedelta64) -> Timedelta: ...
178+
def __sub__(self, other: timedelta | Timedelta | np.timedelta64) -> Self: ...
176179
@overload
177180
def __sub__(self, other: NaTType) -> NaTType: ...
178181
@overload
@@ -182,11 +185,9 @@ class Timedelta(timedelta):
182185
@overload
183186
def __sub__(self, other: pd.TimedeltaIndex) -> TimedeltaIndex: ...
184187
@overload
185-
def __sub__(self, other: Series[pd.Timedelta]) -> Series[pd.Timedelta]: ...
186-
@overload
187-
def __rsub__(self, other: timedelta | Timedelta | np.timedelta64) -> Timedelta: ...
188+
def __rsub__(self, other: timedelta | Timedelta | np.timedelta64) -> Self: ...
188189
@overload
189-
def __rsub__(self, other: dt.datetime | Timestamp | np.datetime64) -> Timestamp: ... # type: ignore[misc]
190+
def __rsub__(self, other: datetime | Timestamp | np.datetime64) -> Timestamp: ... # type: ignore[misc]
190191
@overload
191192
def __rsub__(self, other: NaTType) -> NaTType: ...
192193
@overload
@@ -205,44 +206,31 @@ class Timedelta(timedelta):
205206
) -> np_ndarray[ShapeT, np.timedelta64]: ...
206207
@overload
207208
def __rsub__(self, other: pd.TimedeltaIndex) -> pd.TimedeltaIndex: ...
208-
def __neg__(self) -> Timedelta: ...
209-
def __pos__(self) -> Timedelta: ...
210-
def __abs__(self) -> Timedelta: ...
211-
# Override due to more types supported than dt.timedelta
209+
def __neg__(self) -> Self: ...
210+
def __pos__(self) -> Self: ...
211+
def __abs__(self) -> Self: ...
212+
# Override due to more types supported than timedelta
212213
@overload # type: ignore[override]
213-
def __mul__(self, other: float) -> Timedelta: ...
214+
def __mul__(self, other: float) -> Self: ...
214215
@overload
215216
def __mul__(
216-
self, other: np_ndarray[ShapeT, np.integer] | np_ndarray[ShapeT, np.floating]
217+
self, other: np_ndarray[ShapeT, np.bool_ | np.integer | np.floating]
217218
) -> np_ndarray[ShapeT, np.timedelta64]: ...
218219
@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: ...
224-
@overload
225-
def __rmul__(self, other: float) -> Timedelta: ...
220+
def __rmul__(self, other: float) -> Self: ...
226221
@overload
227222
def __rmul__(
228-
self, other: np_ndarray[ShapeT, np.floating] | np_ndarray[ShapeT, np.integer]
223+
self, other: np_ndarray[ShapeT, np.bool_ | np.integer | np.floating]
229224
) -> 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
225+
# Override due to more types supported than timedelta
238226
# error: Signature of "__floordiv__" incompatible with supertype "timedelta"
239227
@overload # type: ignore[override]
240228
def __floordiv__(self, other: timedelta | Timedelta | np.timedelta64) -> int: ...
241229
@overload
242-
def __floordiv__(self, other: float) -> Timedelta: ...
230+
def __floordiv__(self, other: float) -> Self: ...
243231
@overload
244232
def __floordiv__(
245-
self, other: np_ndarray[ShapeT, np.integer] | np_ndarray[ShapeT, np.floating]
233+
self, other: np_ndarray[ShapeT, np.integer | np.floating]
246234
) -> np_ndarray[ShapeT, np.timedelta64]: ...
247235
@overload
248236
def __floordiv__(
@@ -266,14 +254,14 @@ class Timedelta(timedelta):
266254
def __rfloordiv__(
267255
self, other: np_ndarray[ShapeT, np.timedelta64]
268256
) -> np_ndarray[ShapeT, np.int_]: ...
269-
# Override due to more types supported than dt.timedelta
257+
# Override due to more types supported than timedelta
270258
@overload # type: ignore[override]
271259
def __truediv__(self, other: timedelta | Timedelta | NaTType) -> float: ...
272260
@overload
273-
def __truediv__(self, other: float) -> Timedelta: ...
261+
def __truediv__(self, other: float) -> Self: ...
274262
@overload
275263
def __truediv__(
276-
self, other: np_ndarray[ShapeT, np.integer] | np_ndarray[ShapeT, np.floating]
264+
self, other: np_ndarray[ShapeT, np.integer | np.floating]
277265
) -> np_ndarray[ShapeT, np.timedelta64]: ...
278266
@overload
279267
def __truediv__(self, other: Series[Timedelta]) -> Series[float]: ...
@@ -284,7 +272,7 @@ class Timedelta(timedelta):
284272
@overload
285273
def __truediv__(self, other: Index[int] | Index[float]) -> TimedeltaIndex: ...
286274
def __rtruediv__(self, other: timedelta | Timedelta | NaTType) -> float: ...
287-
# Override due to more types supported than dt.timedelta
275+
# Override due to more types supported than timedelta
288276
@overload
289277
def __eq__(self, other: timedelta | Timedelta | np.timedelta64) -> bool: ... # type: ignore[overload-overlap] # pyright: ignore[reportOverlappingOverload]
290278
@overload
@@ -297,7 +285,7 @@ class Timedelta(timedelta):
297285
) -> np_ndarray[ShapeT, np.bool_]: ...
298286
@overload
299287
def __eq__(self, other: object) -> Literal[False]: ...
300-
# Override due to more types supported than dt.timedelta
288+
# Override due to more types supported than timedelta
301289
@overload
302290
def __ne__(self, other: timedelta | Timedelta | np.timedelta64) -> bool: ... # type: ignore[overload-overlap] # pyright: ignore[reportOverlappingOverload]
303291
@overload
@@ -310,18 +298,18 @@ class Timedelta(timedelta):
310298
) -> np_ndarray[ShapeT, np.bool_]: ...
311299
@overload
312300
def __ne__(self, other: object) -> Literal[True]: ...
313-
# Override due to more types supported than dt.timedelta
301+
# Override due to more types supported than timedelta
314302
@overload # type: ignore[override]
315-
def __mod__(self, other: timedelta) -> Timedelta: ...
303+
def __mod__(self, other: timedelta) -> Self: ...
316304
@overload
317-
def __mod__(self, other: float) -> Timedelta: ...
305+
def __mod__(self, other: float) -> Self: ...
318306
@overload
319307
def __mod__(self, other: Series[int] | Series[float]) -> Series[Timedelta]: ...
320308
@overload
321309
def __mod__(self, other: Index[int] | Index[float]) -> TimedeltaIndex: ...
322310
@overload
323311
def __mod__(
324-
self, other: np_ndarray[ShapeT, np.integer] | np_ndarray[ShapeT, np.floating]
312+
self, other: np_ndarray[ShapeT, np.integer | np.floating]
325313
) -> np_ndarray[ShapeT, np.timedelta64]: ...
326314
@overload
327315
def __mod__(
@@ -330,7 +318,7 @@ class Timedelta(timedelta):
330318
def __divmod__(self, other: timedelta) -> tuple[int, Timedelta]: ...
331319
# Mypy complains Forward operator "<inequality op>" is not callable, so ignore misc
332320
# for le, lt ge and gt
333-
# Override due to more types supported than dt.timedelta
321+
# Override due to more types supported than timedelta
334322
@overload # type: ignore[override]
335323
def __le__(self, other: timedelta | Timedelta | np.timedelta64) -> bool: ... # type: ignore[misc]
336324
@overload
@@ -341,7 +329,7 @@ class Timedelta(timedelta):
341329
) -> np_ndarray[ShapeT, np.bool_]: ...
342330
@overload
343331
def __le__(self, other: Series[pd.Timedelta]) -> Series[bool]: ...
344-
# Override due to more types supported than dt.timedelta
332+
# Override due to more types supported than timedelta
345333
@overload # type: ignore[override]
346334
def __lt__(self, other: timedelta | Timedelta | np.timedelta64) -> bool: ... # type: ignore[misc]
347335
@overload
@@ -352,7 +340,7 @@ class Timedelta(timedelta):
352340
) -> np_ndarray[ShapeT, np.bool_]: ...
353341
@overload
354342
def __lt__(self, other: Series[pd.Timedelta]) -> Series[bool]: ...
355-
# Override due to more types supported than dt.timedelta
343+
# Override due to more types supported than timedelta
356344
@overload # type: ignore[override]
357345
def __ge__(self, other: timedelta | Timedelta | np.timedelta64) -> bool: ... # type: ignore[misc]
358346
@overload
@@ -363,7 +351,7 @@ class Timedelta(timedelta):
363351
) -> np_ndarray[ShapeT, np.bool_]: ...
364352
@overload
365353
def __ge__(self, other: Series[pd.Timedelta]) -> Series[bool]: ...
366-
# Override due to more types supported than dt.timedelta
354+
# Override due to more types supported than timedelta
367355
@overload # type: ignore[override]
368356
def __gt__(self, other: timedelta | Timedelta | np.timedelta64) -> bool: ... # type: ignore[misc]
369357
@overload

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)