Skip to content

Navigation Menu

Sign in
Appearance settings

Search code, repositories, users, issues, pull requests...

Provide feedback

We read every piece of feedback, and take your input very seriously.

Saved searches

Use saved searches to filter your results more quickly

Appearance settings

Commit 17864c5

Browse filesBrowse files
authored
Merge pull request #28266 from charris/backport-28259
TYP: Fix ``timedelta64.__divmod__`` and ``timedelta64.__mod__`` overloads
2 parents feeba15 + caf0748 commit 17864c5
Copy full SHA for 17864c5

File tree

Expand file treeCollapse file tree

2 files changed

+134
-57
lines changed
Filter options
Expand file treeCollapse file tree

2 files changed

+134
-57
lines changed

‎numpy/__init__.pyi

Copy file name to clipboardExpand all lines: numpy/__init__.pyi
+49-8Lines changed: 49 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -4414,10 +4414,12 @@ class timedelta64(_IntegralMixin, generic[_TD64ItemT_co], Generic[_TD64ItemT_co]
44144414
@overload
44154415
def __init__(self: timedelta64[None], value: _NaTValue | None, format: _TimeUnitSpec, /) -> None: ...
44164416
@overload
4417-
def __init__(self: timedelta64[int], value: dt.timedelta, format: _TimeUnitSpec[_IntTimeUnit], /) -> None: ...
4417+
def __init__(self: timedelta64[L[0]], value: L[0], format: _TimeUnitSpec[_IntTD64Unit] = ..., /) -> None: ...
44184418
@overload
44194419
def __init__(self: timedelta64[int], value: _IntLike_co, format: _TimeUnitSpec[_IntTD64Unit] = ..., /) -> None: ...
44204420
@overload
4421+
def __init__(self: timedelta64[int], value: dt.timedelta, format: _TimeUnitSpec[_IntTimeUnit], /) -> None: ...
4422+
@overload
44214423
def __init__(
44224424
self: timedelta64[dt.timedelta],
44234425
value: dt.timedelta | _IntLike_co,
@@ -4458,29 +4460,68 @@ class timedelta64(_IntegralMixin, generic[_TD64ItemT_co], Generic[_TD64ItemT_co]
44584460
def __mul__(self, x: float | np.floating[Any] | np.integer[Any] | np.bool, /) -> timedelta64: ...
44594461
__rmul__ = __mul__
44604462

4463+
@overload
4464+
def __mod__(self, x: timedelta64[None | L[0]], /) -> timedelta64[None]: ...
44614465
@overload
44624466
def __mod__(self: timedelta64[None], x: timedelta64, /) -> timedelta64[None]: ...
44634467
@overload
4468+
def __mod__(self: timedelta64[int], x: timedelta64[int | dt.timedelta], /) -> timedelta64[int | None]: ...
4469+
@overload
4470+
def __mod__(self: timedelta64[dt.timedelta], x: timedelta64[_AnyTD64Item], /) -> timedelta64[_AnyTD64Item | None]: ...
4471+
@overload
44644472
def __mod__(self: timedelta64[dt.timedelta], x: dt.timedelta, /) -> dt.timedelta: ...
44654473
@overload
4466-
def __mod__(self: timedelta64[dt.timedelta], x: timedelta64[_AnyTD64Item], /) -> timedelta64[_AnyTD64Item]: ...
4474+
def __mod__(self, x: timedelta64[int], /) -> timedelta64[int | None]: ...
44674475
@overload
4468-
def __mod__(self: timedelta64[int], x: timedelta64[int | dt.timedelta], /) -> timedelta64[int]: ...
4476+
def __mod__(self, x: timedelta64, /) -> timedelta64: ...
4477+
4478+
# the L[0] makes __mod__ non-commutative, which the first two overloads reflect
44694479
@overload
4470-
def __mod__(self, x: timedelta64[None], /) -> timedelta64[None]: ...
4480+
def __rmod__(self, x: timedelta64[None], /) -> timedelta64[None]: ...
44714481
@overload
4472-
def __mod__(self, x: timedelta64[int], /) -> timedelta64[int]: ...
4482+
def __rmod__(self: timedelta64[None | L[0]], x: timedelta64, /) -> timedelta64[None]: ...
44734483
@overload
4474-
def __mod__(self, x: timedelta64, /) -> timedelta64: ...
4475-
__rmod__ = __mod__ # at runtime the outcomes differ, but the type signatures are the same
4484+
def __rmod__(self: timedelta64[int], x: timedelta64[int | dt.timedelta], /) -> timedelta64[int | None]: ...
4485+
@overload
4486+
def __rmod__(self: timedelta64[dt.timedelta], x: timedelta64[_AnyTD64Item], /) -> timedelta64[_AnyTD64Item | None]: ...
4487+
@overload
4488+
def __rmod__(self: timedelta64[dt.timedelta], x: dt.timedelta, /) -> dt.timedelta: ...
4489+
@overload
4490+
def __rmod__(self, x: timedelta64[int], /) -> timedelta64[int | None]: ...
4491+
@overload
4492+
def __rmod__(self, x: timedelta64, /) -> timedelta64: ...
44764493

4494+
# keep in sync with __mod__
4495+
@overload
4496+
def __divmod__(self, x: timedelta64[None | L[0]], /) -> tuple[int64, timedelta64[None]]: ...
44774497
@overload
44784498
def __divmod__(self: timedelta64[None], x: timedelta64, /) -> tuple[int64, timedelta64[None]]: ...
44794499
@overload
4500+
def __divmod__(self: timedelta64[int], x: timedelta64[int | dt.timedelta], /) -> tuple[int64, timedelta64[int | None]]: ...
4501+
@overload
4502+
def __divmod__(self: timedelta64[dt.timedelta], x: timedelta64[_AnyTD64Item], /) -> tuple[int64, timedelta64[_AnyTD64Item | None]]: ...
4503+
@overload
44804504
def __divmod__(self: timedelta64[dt.timedelta], x: dt.timedelta, /) -> tuple[int, dt.timedelta]: ...
44814505
@overload
4506+
def __divmod__(self, x: timedelta64[int], /) -> tuple[int64, timedelta64[int | None]]: ...
4507+
@overload
44824508
def __divmod__(self, x: timedelta64, /) -> tuple[int64, timedelta64]: ...
4483-
__rdivmod__ = __divmod__
4509+
4510+
# keep in sync with __rmod__
4511+
@overload
4512+
def __rdivmod__(self, x: timedelta64[None], /) -> tuple[int64, timedelta64[None]]: ...
4513+
@overload
4514+
def __rdivmod__(self: timedelta64[None | L[0]], x: timedelta64, /) -> tuple[int64, timedelta64[None]]: ...
4515+
@overload
4516+
def __rdivmod__(self: timedelta64[int], x: timedelta64[int | dt.timedelta], /) -> tuple[int64, timedelta64[int | None]]: ...
4517+
@overload
4518+
def __rdivmod__(self: timedelta64[dt.timedelta], x: timedelta64[_AnyTD64Item], /) -> tuple[int64, timedelta64[_AnyTD64Item | None]]: ...
4519+
@overload
4520+
def __rdivmod__(self: timedelta64[dt.timedelta], x: dt.timedelta, /) -> tuple[int, dt.timedelta]: ...
4521+
@overload
4522+
def __rdivmod__(self, x: timedelta64[int], /) -> tuple[int64, timedelta64[int | None]]: ...
4523+
@overload
4524+
def __rdivmod__(self, x: timedelta64, /) -> tuple[int64, timedelta64]: ...
44844525

44854526
@overload
44864527
def __sub__(self: timedelta64[None], b: _TD64Like_co, /) -> timedelta64[None]: ...
+85-49Lines changed: 85 additions & 49 deletions
Original file line numberDiff line numberDiff line change
@@ -1,39 +1,71 @@
11
import datetime as dt
2-
from typing import Any
2+
from typing import Literal as L
3+
4+
from typing_extensions import assert_type
35

46
import numpy as np
57
import numpy.typing as npt
6-
from numpy._typing import _32Bit, _64Bit
8+
from numpy._typing import _64Bit
79

8-
from typing_extensions import assert_type
10+
f8: np.float64
11+
i8: np.int64
12+
u8: np.uint64
913

10-
f8 = np.float64()
11-
i8 = np.int64()
12-
u8 = np.uint64()
14+
f4: np.float32
15+
i4: np.int32
16+
u4: np.uint32
1317

14-
f4 = np.float32()
15-
i4 = np.int32()
16-
u4 = np.uint32()
18+
m: np.timedelta64
19+
m_nat: np.timedelta64[None]
20+
m_int0: np.timedelta64[L[0]]
21+
m_int: np.timedelta64[int]
22+
m_td: np.timedelta64[dt.timedelta]
1723

18-
td = np.timedelta64(0, "D")
19-
b_ = np.bool()
24+
b_: np.bool
2025

21-
b = bool()
22-
f = float()
23-
i = int()
26+
b: bool
27+
i: int
28+
f: float
2429

2530
AR_b: npt.NDArray[np.bool]
2631
AR_m: npt.NDArray[np.timedelta64]
2732

2833
# Time structures
2934

30-
assert_type(td % td, np.timedelta64[dt.timedelta])
31-
assert_type(AR_m % td, npt.NDArray[np.timedelta64])
32-
assert_type(td % AR_m, npt.NDArray[np.timedelta64])
33-
34-
assert_type(divmod(td, td), tuple[np.int64, np.timedelta64])
35-
assert_type(divmod(AR_m, td), tuple[npt.NDArray[np.int64], npt.NDArray[np.timedelta64]])
36-
assert_type(divmod(td, AR_m), tuple[npt.NDArray[np.int64], npt.NDArray[np.timedelta64]])
35+
assert_type(m % m, np.timedelta64)
36+
assert_type(m % m_nat, np.timedelta64[None])
37+
assert_type(m % m_int0, np.timedelta64[None])
38+
assert_type(m % m_int, np.timedelta64[int | None])
39+
assert_type(m_nat % m, np.timedelta64[None])
40+
assert_type(m_int % m_nat, np.timedelta64[None])
41+
assert_type(m_int % m_int0, np.timedelta64[None])
42+
assert_type(m_int % m_int, np.timedelta64[int | None])
43+
assert_type(m_int % m_td, np.timedelta64[int | None])
44+
assert_type(m_td % m_nat, np.timedelta64[None])
45+
assert_type(m_td % m_int0, np.timedelta64[None])
46+
assert_type(m_td % m_int, np.timedelta64[int | None])
47+
assert_type(m_td % m_td, np.timedelta64[dt.timedelta | None])
48+
49+
assert_type(AR_m % m, npt.NDArray[np.timedelta64])
50+
assert_type(m % AR_m, npt.NDArray[np.timedelta64])
51+
52+
assert_type(divmod(m, m), tuple[np.int64, np.timedelta64])
53+
assert_type(divmod(m, m_nat), tuple[np.int64, np.timedelta64[None]])
54+
assert_type(divmod(m, m_int0), tuple[np.int64, np.timedelta64[None]])
55+
# workarounds for https://github.com/microsoft/pyright/issues/9663
56+
assert_type(m.__divmod__(m_int), tuple[np.int64, np.timedelta64[int | None]])
57+
assert_type(divmod(m_nat, m), tuple[np.int64, np.timedelta64[None]])
58+
assert_type(divmod(m_int, m_nat), tuple[np.int64, np.timedelta64[None]])
59+
assert_type(divmod(m_int, m_int0), tuple[np.int64, np.timedelta64[None]])
60+
assert_type(divmod(m_int, m_int), tuple[np.int64, np.timedelta64[int | None]])
61+
assert_type(divmod(m_int, m_td), tuple[np.int64, np.timedelta64[int | None]])
62+
assert_type(divmod(m_td, m_nat), tuple[np.int64, np.timedelta64[None]])
63+
assert_type(divmod(m_td, m_int0), tuple[np.int64, np.timedelta64[None]])
64+
assert_type(divmod(m_td, m_int), tuple[np.int64, np.timedelta64[int | None]])
65+
assert_type(divmod(m_td, m_td), tuple[np.int64, np.timedelta64[dt.timedelta | None]])
66+
67+
assert_type(divmod(AR_m, m), tuple[npt.NDArray[np.int64], npt.NDArray[np.timedelta64]])
68+
assert_type(divmod(m, AR_m), tuple[npt.NDArray[np.int64], npt.NDArray[np.timedelta64]])
3769

3870
# Bool
3971

@@ -47,11 +79,12 @@ assert_type(b_ % f8, np.float64)
4779
assert_type(b_ % AR_b, npt.NDArray[np.int8])
4880

4981
assert_type(divmod(b_, b), tuple[np.int8, np.int8])
50-
assert_type(divmod(b_, i), tuple[np.int_, np.int_])
51-
assert_type(divmod(b_, f), tuple[np.float64, np.float64])
5282
assert_type(divmod(b_, b_), tuple[np.int8, np.int8])
53-
assert_type(divmod(b_, i8), tuple[np.int64, np.int64])
54-
assert_type(divmod(b_, u8), tuple[np.uint64, np.uint64])
83+
# workarounds for https://github.com/microsoft/pyright/issues/9663
84+
assert_type(b_.__divmod__(i), tuple[np.int_, np.int_])
85+
assert_type(b_.__divmod__(f), tuple[np.float64, np.float64])
86+
assert_type(b_.__divmod__(i8), tuple[np.int64, np.int64])
87+
assert_type(b_.__divmod__(u8), tuple[np.uint64, np.uint64])
5588
assert_type(divmod(b_, f8), tuple[np.float64, np.float64])
5689
assert_type(divmod(b_, AR_b), tuple[npt.NDArray[np.int8], npt.NDArray[np.int8]])
5790

@@ -77,26 +110,27 @@ assert_type(divmod(AR_b, b_), tuple[npt.NDArray[np.int8], npt.NDArray[np.int8]])
77110

78111
assert_type(i8 % b, np.int64)
79112
assert_type(i8 % i8, np.int64)
80-
assert_type(i8 % f, np.floating[_64Bit])
81-
assert_type(i8 % f8, np.floating[_64Bit])
113+
assert_type(i8 % f, np.float64 | np.floating[_64Bit])
114+
assert_type(i8 % f8, np.float64 | np.floating[_64Bit])
82115
assert_type(i4 % i8, np.int64 | np.int32)
83116
assert_type(i4 % f8, np.float64 | np.float32)
84117
assert_type(i4 % i4, np.int32)
85118
assert_type(i4 % f4, np.float32)
86119
assert_type(i8 % AR_b, npt.NDArray[np.int64])
87120

88-
assert_type(divmod(i8, b), tuple[np.signedinteger[_64Bit], np.signedinteger[_64Bit]])
89-
assert_type(divmod(i8, f), tuple[np.floating[_64Bit], np.floating[_64Bit]])
90-
assert_type(divmod(i8, i8), tuple[np.signedinteger[_64Bit], np.signedinteger[_64Bit]])
91-
assert_type(divmod(i8, f8), tuple[np.floating[_64Bit], np.floating[_64Bit]])
92-
assert_type(divmod(i8, i4), tuple[np.signedinteger[_64Bit], np.signedinteger[_64Bit]] | tuple[np.signedinteger[_32Bit], np.signedinteger[_32Bit]])
93-
assert_type(divmod(i8, f4), tuple[np.floating[_64Bit], np.floating[_64Bit]] | tuple[np.floating[_32Bit], np.floating[_32Bit]])
94-
assert_type(divmod(i4, i4), tuple[np.signedinteger[_32Bit], np.signedinteger[_32Bit]])
95-
assert_type(divmod(i4, f4), tuple[np.floating[_32Bit], np.floating[_32Bit]])
121+
assert_type(divmod(i8, b), tuple[np.int64, np.int64])
122+
assert_type(divmod(i8, i4), tuple[np.int64, np.int64] | tuple[np.int32, np.int32])
123+
assert_type(divmod(i8, i8), tuple[np.int64, np.int64])
124+
# workarounds for https://github.com/microsoft/pyright/issues/9663
125+
assert_type(i8.__divmod__(f), tuple[np.floating[_64Bit], np.floating[_64Bit]] | tuple[np.float64, np.float64])
126+
assert_type(i8.__divmod__(f8), tuple[np.floating[_64Bit], np.floating[_64Bit]] | tuple[np.float64, np.float64])
127+
assert_type(divmod(i8, f4), tuple[np.floating[_64Bit], np.floating[_64Bit]] | tuple[np.float32, np.float32])
128+
assert_type(divmod(i4, i4), tuple[np.int32, np.int32])
129+
assert_type(divmod(i4, f4), tuple[np.float32, np.float32])
96130
assert_type(divmod(i8, AR_b), tuple[npt.NDArray[np.int64], npt.NDArray[np.int64]])
97131

98-
assert_type(b % i8, np.signedinteger[_64Bit])
99-
assert_type(f % i8, np.floating[_64Bit])
132+
assert_type(b % i8, np.int64)
133+
assert_type(f % i8, np.float64 | np.floating[_64Bit])
100134
assert_type(i8 % i8, np.int64)
101135
assert_type(f8 % i8, np.float64)
102136
assert_type(i8 % i4, np.int64 | np.int32)
@@ -105,21 +139,22 @@ assert_type(i4 % i4, np.int32)
105139
assert_type(f4 % i4, np.float32)
106140
assert_type(AR_b % i8, npt.NDArray[np.int64])
107141

108-
assert_type(divmod(b, i8), tuple[np.signedinteger[_64Bit], np.signedinteger[_64Bit]])
109-
assert_type(divmod(f, i8), tuple[np.floating[_64Bit], np.floating[_64Bit]])
142+
assert_type(divmod(b, i8), tuple[np.int64, np.int64])
143+
assert_type(divmod(f, i8), tuple[np.floating[_64Bit], np.floating[_64Bit]] | tuple[np.float64, np.float64])
110144
assert_type(divmod(i8, i8), tuple[np.int64, np.int64])
111145
assert_type(divmod(f8, i8), tuple[np.float64, np.float64])
112-
assert_type(divmod(i4, i8), tuple[np.signedinteger[_64Bit], np.signedinteger[_64Bit]] | tuple[np.signedinteger[_32Bit], np.signedinteger[_32Bit]])
113-
assert_type(divmod(f4, i8), tuple[np.floating[_64Bit], np.floating[_64Bit]] | tuple[np.floating[_32Bit], np.floating[_32Bit]])
114-
assert_type(divmod(i4, i4), tuple[np.signedinteger[_32Bit], np.signedinteger[_32Bit]])
115-
assert_type(divmod(f4, i4), tuple[np.floating[_32Bit], np.floating[_32Bit]])
116-
assert_type(divmod(AR_b, i8), tuple[npt.NDArray[np.int64], npt.NDArray[np.int64]])
146+
assert_type(divmod(i4, i8), tuple[np.int64, np.int64] | tuple[np.int32, np.int32])
147+
assert_type(divmod(i4, i4), tuple[np.int32, np.int32])
148+
# workarounds for https://github.com/microsoft/pyright/issues/9663
149+
assert_type(f4.__divmod__(i8), tuple[np.floating[_64Bit], np.floating[_64Bit]] | tuple[np.float32, np.float32])
150+
assert_type(f4.__divmod__(i4), tuple[np.float32, np.float32])
151+
assert_type(AR_b.__divmod__(i8), tuple[npt.NDArray[np.int64], npt.NDArray[np.int64]])
117152

118153
# float
119154

120155
assert_type(f8 % b, np.float64)
121156
assert_type(f8 % f, np.float64)
122-
assert_type(i8 % f4, np.floating[_64Bit] | np.floating[_32Bit])
157+
assert_type(i8 % f4, np.floating[_64Bit] | np.float32)
123158
assert_type(f4 % f4, np.float32)
124159
assert_type(f8 % AR_b, npt.NDArray[np.float64])
125160

@@ -131,15 +166,16 @@ assert_type(divmod(f4, f4), tuple[np.float32, np.float32])
131166
assert_type(divmod(f8, AR_b), tuple[npt.NDArray[np.float64], npt.NDArray[np.float64]])
132167

133168
assert_type(b % f8, np.float64)
134-
assert_type(f % f8, np.float64)
169+
assert_type(f % f8, np.float64) # pyright: ignore[reportAssertTypeFailure] # pyright incorrectly infers `builtins.float`
135170
assert_type(f8 % f8, np.float64)
136171
assert_type(f8 % f8, np.float64)
137172
assert_type(f4 % f4, np.float32)
138173
assert_type(AR_b % f8, npt.NDArray[np.float64])
139174

140175
assert_type(divmod(b, f8), tuple[np.float64, np.float64])
141-
assert_type(divmod(f, f8), tuple[np.float64, np.float64])
142176
assert_type(divmod(f8, f8), tuple[np.float64, np.float64])
143-
assert_type(divmod(f4, f8), tuple[np.float64, np.float64] | tuple[np.float32, np.float32])
144177
assert_type(divmod(f4, f4), tuple[np.float32, np.float32])
145-
assert_type(divmod(AR_b, f8), tuple[npt.NDArray[np.float64], npt.NDArray[np.float64]])
178+
# workarounds for https://github.com/microsoft/pyright/issues/9663
179+
assert_type(f8.__rdivmod__(f), tuple[np.float64, np.float64])
180+
assert_type(f8.__rdivmod__(f4), tuple[np.float64, np.float64])
181+
assert_type(AR_b.__divmod__(f8), tuple[npt.NDArray[np.float64], npt.NDArray[np.float64]])

0 commit comments

Comments
0 (0)
Morty Proxy This is a proxified and sanitized view of the page, visit original site.