Skip to content

Commit 17864c5

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

2 files changed

Lines changed: 134 additions & 57 deletions

File tree

‎numpy/__init__.pyi‎

Lines 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]: ...

‎numpy/typing/tests/data/reveal/mod.pyi‎

Lines 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)