codekingpro/portable-devtools

1import sys
2from collections import deque
3from pathlib import Path
4from typing import Any, Generic, TypeVar, assert_type
5
6import numpy as np
7import numpy.typing as npt
8
9_ScalarT_co = TypeVar("_ScalarT_co", bound=np.generic, covariant=True)
10
11class SubClass(npt.NDArray[_ScalarT_co]): ...
12
13class IntoSubClass(Generic[_ScalarT_co]):
14    def __array__(self) -> SubClass[_ScalarT_co]: ...
15
16i8: np.int64
17
18A: npt.NDArray[np.float64]
19B: SubClass[np.float64]
20C: list[int]
21D: SubClass[np.float64 | np.int64]
22E: IntoSubClass[np.float64 | np.int64]
23
24mixed_shape: tuple[int, np.int64]
25
26def func(i: int, j: int, **kwargs: Any) -> SubClass[np.float64]: ...
27
28assert_type(np.empty_like(A), npt.NDArray[np.float64])
29assert_type(np.empty_like(B), SubClass[np.float64])
30assert_type(np.empty_like([1, 1.0]), npt.NDArray[Any])
31assert_type(np.empty_like(A, dtype=np.int64), npt.NDArray[np.int64])
32assert_type(np.empty_like(A, dtype="c16"), npt.NDArray[Any])
33
34assert_type(np.array(A), npt.NDArray[np.float64])
35assert_type(np.array(B), npt.NDArray[np.float64])
36assert_type(np.array([1, 1.0]), npt.NDArray[Any])
37assert_type(np.array(deque([1, 2, 3])), npt.NDArray[Any])
38assert_type(np.array(A, dtype=np.int64), npt.NDArray[np.int64])
39assert_type(np.array(A, dtype="c16"), npt.NDArray[Any])
40assert_type(np.array(A, like=A), npt.NDArray[np.float64])
41assert_type(np.array(A, subok=True), npt.NDArray[np.float64])
42assert_type(np.array(B, subok=True), SubClass[np.float64])
43assert_type(np.array(B, subok=True, ndmin=0), SubClass[np.float64])
44assert_type(np.array(B, subok=True, ndmin=1), SubClass[np.float64])
45assert_type(np.array(D), npt.NDArray[np.float64 | np.int64])
46assert_type(np.array(E, subok=True), SubClass[np.float64 | np.int64])
47# https://github.com/numpy/numpy/issues/29245
48assert_type(np.array([], dtype=np.bool), npt.NDArray[np.bool])
49
50assert_type(np.zeros([1, 5, 6]), npt.NDArray[np.float64])
51assert_type(np.zeros([1, 5, 6], dtype=np.int64), npt.NDArray[np.int64])
52assert_type(np.zeros([1, 5, 6], dtype="c16"), npt.NDArray[Any])
53assert_type(np.zeros(mixed_shape), npt.NDArray[np.float64])
54
55assert_type(np.empty([1, 5, 6]), npt.NDArray[np.float64])
56assert_type(np.empty([1, 5, 6], dtype=np.int64), npt.NDArray[np.int64])
57assert_type(np.empty([1, 5, 6], dtype="c16"), npt.NDArray[Any])
58assert_type(np.empty(mixed_shape), npt.NDArray[np.float64])
59
60assert_type(np.concatenate(A), npt.NDArray[np.float64])
61assert_type(np.concatenate([A, A]), npt.NDArray[Any])  # pyright correctly infers this as NDArray[float64]
62assert_type(np.concatenate([[1], A]), npt.NDArray[Any])
63assert_type(np.concatenate([[1], [1]]), npt.NDArray[Any])
64assert_type(np.concatenate((A, A)), npt.NDArray[np.float64])
65assert_type(np.concatenate(([1], [1])), npt.NDArray[Any])
66assert_type(np.concatenate([1, 1.0]), npt.NDArray[Any])
67assert_type(np.concatenate(A, dtype=np.int64), npt.NDArray[np.int64])
68assert_type(np.concatenate(A, dtype="c16"), npt.NDArray[Any])
69assert_type(np.concatenate([1, 1.0], out=A), npt.NDArray[np.float64])
70
71assert_type(np.asarray(A), npt.NDArray[np.float64])
72assert_type(np.asarray(B), npt.NDArray[np.float64])
73assert_type(np.asarray([1, 1.0]), npt.NDArray[Any])
74assert_type(np.asarray(A, dtype=np.int64), npt.NDArray[np.int64])
75assert_type(np.asarray(A, dtype="c16"), npt.NDArray[Any])
76
77assert_type(np.asanyarray(A), npt.NDArray[np.float64])
78assert_type(np.asanyarray(B), SubClass[np.float64])
79assert_type(np.asanyarray([1, 1.0]), npt.NDArray[Any])
80assert_type(np.asanyarray(A, dtype=np.int64), npt.NDArray[np.int64])
81assert_type(np.asanyarray(A, dtype="c16"), npt.NDArray[Any])
82
83assert_type(np.ascontiguousarray(A), npt.NDArray[np.float64])
84assert_type(np.ascontiguousarray(B), npt.NDArray[np.float64])
85assert_type(np.ascontiguousarray([1, 1.0]), npt.NDArray[Any])
86assert_type(np.ascontiguousarray(A, dtype=np.int64), npt.NDArray[np.int64])
87assert_type(np.ascontiguousarray(A, dtype="c16"), npt.NDArray[Any])
88
89assert_type(np.asfortranarray(A), npt.NDArray[np.float64])
90assert_type(np.asfortranarray(B), npt.NDArray[np.float64])
91assert_type(np.asfortranarray([1, 1.0]), npt.NDArray[Any])
92assert_type(np.asfortranarray(A, dtype=np.int64), npt.NDArray[np.int64])
93assert_type(np.asfortranarray(A, dtype="c16"), npt.NDArray[Any])
94
95assert_type(np.fromstring("1 1 1", sep=" "), npt.NDArray[np.float64])
96assert_type(np.fromstring(b"1 1 1", sep=" "), npt.NDArray[np.float64])
97assert_type(np.fromstring("1 1 1", dtype=np.int64, sep=" "), npt.NDArray[np.int64])
98assert_type(np.fromstring(b"1 1 1", dtype=np.int64, sep=" "), npt.NDArray[np.int64])
99assert_type(np.fromstring("1 1 1", dtype="c16", sep=" "), npt.NDArray[Any])
100assert_type(np.fromstring(b"1 1 1", dtype="c16", sep=" "), npt.NDArray[Any])
101
102assert_type(np.fromfile("test.txt", sep=" "), npt.NDArray[np.float64])
103assert_type(np.fromfile("test.txt", dtype=np.int64, sep=" "), npt.NDArray[np.int64])
104assert_type(np.fromfile("test.txt", dtype="c16", sep=" "), npt.NDArray[Any])
105with open("test.txt") as f:
106    assert_type(np.fromfile(f, sep=" "), npt.NDArray[np.float64])
107    assert_type(np.fromfile(b"test.txt", sep=" "), npt.NDArray[np.float64])
108    assert_type(np.fromfile(Path("test.txt"), sep=" "), npt.NDArray[np.float64])
109
110assert_type(np.fromiter("12345", np.float64), npt.NDArray[np.float64])
111assert_type(np.fromiter("12345", float), npt.NDArray[Any])
112
113assert_type(np.frombuffer(A), npt.NDArray[np.float64])
114assert_type(np.frombuffer(A, dtype=np.int64), npt.NDArray[np.int64])
115assert_type(np.frombuffer(A, dtype="c16"), npt.NDArray[Any])
116
117_x_bool: bool
118_x_int: int
119_x_float: float
120_x_timedelta: np.timedelta64
121_x_datetime: np.datetime64
122
123assert_type(np.arange(False, True), np.ndarray[tuple[int], np.dtype[np.int_]])
124assert_type(np.arange(10), np.ndarray[tuple[int], np.dtype[np.int_]])
125assert_type(np.arange(0, 10, step=2), np.ndarray[tuple[int], np.dtype[np.int_]])
126assert_type(np.arange(10.0), np.ndarray[tuple[int], np.dtype[np.float64 | Any]])
127assert_type(np.arange(0, stop=10.0), np.ndarray[tuple[int], np.dtype[np.float64 | Any]])
128assert_type(np.arange(_x_timedelta), np.ndarray[tuple[int], np.dtype[np.timedelta64[Any]]])
129assert_type(np.arange(0, _x_timedelta), np.ndarray[tuple[int], np.dtype[np.timedelta64[Any]]])
130assert_type(np.arange(_x_datetime, _x_datetime), np.ndarray[tuple[int], np.dtype[np.datetime64[Any]]])
131assert_type(np.arange(10, dtype=np.float64), np.ndarray[tuple[int], np.dtype[np.float64]])
132assert_type(np.arange(0, 10, step=2, dtype=np.int16), np.ndarray[tuple[int], np.dtype[np.int16]])
133assert_type(np.arange(10, dtype=int), np.ndarray[tuple[int], np.dtype[np.int_]])
134assert_type(np.arange(0, 10, dtype="f8"), np.ndarray[tuple[int], np.dtype])
135# https://github.com/numpy/numpy/issues/30628
136assert_type(np.arange("2025-12-20", "2025-12-23", dtype="datetime64[D]"), np.ndarray[tuple[int], np.dtype[np.datetime64[Any]]])
137
138assert_type(np.require(A), npt.NDArray[np.float64])
139assert_type(np.require(B), SubClass[np.float64])
140assert_type(np.require(B, requirements=None), SubClass[np.float64])
141assert_type(np.require(B, dtype=int), npt.NDArray[Any])
142assert_type(np.require(B, requirements="E"), npt.NDArray[Any])
143assert_type(np.require(B, requirements=["ENSUREARRAY"]), npt.NDArray[Any])
144assert_type(np.require(B, requirements={"F", "E"}), npt.NDArray[Any])
145assert_type(np.require(B, requirements=["C", "OWNDATA"]), SubClass[np.float64])
146assert_type(np.require(B, requirements="W"), SubClass[np.float64])
147assert_type(np.require(B, requirements="A"), SubClass[np.float64])
148assert_type(np.require(C), npt.NDArray[Any])
149
150assert_type(np.linspace(0, 10), npt.NDArray[np.float64])
151assert_type(np.linspace(0, 10j), npt.NDArray[np.complexfloating])
152assert_type(np.linspace(0, 10, dtype=np.int64), npt.NDArray[np.int64])
153assert_type(np.linspace(0, 10, dtype=int), npt.NDArray[Any])
154assert_type(np.linspace(0, 10, retstep=True), tuple[npt.NDArray[np.float64], np.float64])
155assert_type(np.linspace(0j, 10, retstep=True), tuple[npt.NDArray[np.complexfloating], np.complexfloating])
156assert_type(np.linspace(0, 10, retstep=True, dtype=np.int64), tuple[npt.NDArray[np.int64], np.int64])
157assert_type(np.linspace(0j, 10, retstep=True, dtype=int), tuple[npt.NDArray[Any], Any])
158
159assert_type(np.logspace(0, 10), npt.NDArray[np.float64])
160assert_type(np.logspace(0, 10j), npt.NDArray[np.complexfloating])
161assert_type(np.logspace(0, 10, dtype=np.int64), npt.NDArray[np.int64])
162assert_type(np.logspace(0, 10, dtype=int), npt.NDArray[Any])
163
164assert_type(np.geomspace(0, 10), npt.NDArray[np.float64])
165assert_type(np.geomspace(0, 10j), npt.NDArray[np.complexfloating])
166assert_type(np.geomspace(0, 10, dtype=np.int64), npt.NDArray[np.int64])
167assert_type(np.geomspace(0, 10, dtype=int), npt.NDArray[Any])
168
169assert_type(np.zeros_like(A), npt.NDArray[np.float64])
170assert_type(np.zeros_like(C), npt.NDArray[Any])
171assert_type(np.zeros_like(A, dtype=float), npt.NDArray[Any])
172assert_type(np.zeros_like(B), SubClass[np.float64])
173assert_type(np.zeros_like(B, dtype=np.int64), npt.NDArray[np.int64])
174
175assert_type(np.ones_like(A), npt.NDArray[np.float64])
176assert_type(np.ones_like(C), npt.NDArray[Any])
177assert_type(np.ones_like(A, dtype=float), npt.NDArray[Any])
178assert_type(np.ones_like(B), SubClass[np.float64])
179assert_type(np.ones_like(B, dtype=np.int64), npt.NDArray[np.int64])
180
181assert_type(np.full_like(A, i8), npt.NDArray[np.float64])
182assert_type(np.full_like(C, i8), npt.NDArray[Any])
183assert_type(np.full_like(A, i8, dtype=int), npt.NDArray[Any])
184assert_type(np.full_like(B, i8), SubClass[np.float64])
185assert_type(np.full_like(B, i8, dtype=np.int64), npt.NDArray[np.int64])
186
187_size: int
188_shape_0d: tuple[()]
189_shape_1d: tuple[int]
190_shape_2d: tuple[int, int]
191_shape_nd: tuple[int, ...]
192_shape_like: list[int]
193
194assert_type(np.ones(_shape_0d), np.ndarray[tuple[()], np.dtype[np.float64]])
195assert_type(np.ones(_size), np.ndarray[tuple[int], np.dtype[np.float64]])
196assert_type(np.ones(_shape_2d), np.ndarray[tuple[int, int], np.dtype[np.float64]])
197assert_type(np.ones(_shape_nd), np.ndarray[tuple[int, ...], np.dtype[np.float64]])
198assert_type(np.ones(_shape_1d, dtype=np.int64), np.ndarray[tuple[int], np.dtype[np.int64]])
199assert_type(np.ones(_shape_like), npt.NDArray[np.float64])
200assert_type(
201    np.ones(_shape_like, dtype=np.dtypes.Int64DType()),
202    np.ndarray[tuple[Any, ...], np.dtypes.Int64DType],
203)
204assert_type(np.ones(_shape_like, dtype=int), npt.NDArray[Any])
205assert_type(np.ones(mixed_shape), npt.NDArray[np.float64])
206
207assert_type(np.full(_size, i8), np.ndarray[tuple[int], np.dtype[np.int64]])
208assert_type(np.full(_shape_2d, i8), np.ndarray[tuple[int, int], np.dtype[np.int64]])
209assert_type(np.full(_shape_like, i8), npt.NDArray[np.int64])
210assert_type(np.full(_shape_like, 42), npt.NDArray[Any])
211assert_type(np.full(_size, i8, dtype=np.float64), np.ndarray[tuple[int], np.dtype[np.float64]])
212assert_type(np.full(_size, i8, dtype=float), np.ndarray[tuple[int], np.dtype])
213assert_type(np.full(_shape_like, 42, dtype=float), npt.NDArray[Any])
214assert_type(np.full(_shape_0d, i8, dtype=object), np.ndarray[tuple[()], np.dtype])
215
216assert_type(np.indices([1, 2, 3]), npt.NDArray[np.int_])
217assert_type(np.indices([1, 2, 3], sparse=True), tuple[npt.NDArray[np.int_], ...])
218
219assert_type(np.fromfunction(func, (3, 5)), SubClass[np.float64])
220
221assert_type(np.identity(3), np.ndarray[tuple[int, int], np.dtype[np.float64]])
222assert_type(np.identity(3, dtype=np.int8), np.ndarray[tuple[int, int], np.dtype[np.int8]])
223assert_type(np.identity(3, dtype=bool), np.ndarray[tuple[int, int], np.dtype[np.bool]])
224assert_type(np.identity(3, dtype="bool"), np.ndarray[tuple[int, int], np.dtype[np.bool]])
225assert_type(np.identity(3, dtype="b1"), np.ndarray[tuple[int, int], np.dtype[np.bool]])
226assert_type(np.identity(3, dtype="?"), np.ndarray[tuple[int, int], np.dtype[np.bool]])
227assert_type(np.identity(3, dtype=int), np.ndarray[tuple[int, int], np.dtype[np.int_ | Any]])
228assert_type(np.identity(3, dtype="int"), np.ndarray[tuple[int, int], np.dtype[np.int_ | Any]])
229assert_type(np.identity(3, dtype="n"), np.ndarray[tuple[int, int], np.dtype[np.int_ | Any]])
230assert_type(np.identity(3, dtype=float), np.ndarray[tuple[int, int], np.dtype[np.float64 | Any]])
231assert_type(np.identity(3, dtype="float"), np.ndarray[tuple[int, int], np.dtype[np.float64 | Any]])
232assert_type(np.identity(3, dtype="f8"), np.ndarray[tuple[int, int], np.dtype[np.float64 | Any]])
233assert_type(np.identity(3, dtype="d"), np.ndarray[tuple[int, int], np.dtype[np.float64 | Any]])
234assert_type(np.identity(3, dtype=complex), np.ndarray[tuple[int, int], np.dtype[np.complex128 | Any]])
235assert_type(np.identity(3, dtype="complex"), np.ndarray[tuple[int, int], np.dtype[np.complex128 | Any]])
236assert_type(np.identity(3, dtype="c16"), np.ndarray[tuple[int, int], np.dtype[np.complex128 | Any]])
237assert_type(np.identity(3, dtype="D"), np.ndarray[tuple[int, int], np.dtype[np.complex128 | Any]])
238
239assert_type(np.atleast_1d(A), npt.NDArray[np.float64])
240assert_type(np.atleast_1d(C), npt.NDArray[Any])
241assert_type(np.atleast_1d(A, A), tuple[npt.NDArray[np.float64], npt.NDArray[np.float64]])
242assert_type(np.atleast_1d(A, C), tuple[npt.NDArray[Any], npt.NDArray[Any]])
243assert_type(np.atleast_1d(C, C), tuple[npt.NDArray[Any], npt.NDArray[Any]])
244assert_type(np.atleast_1d(A, A, A), tuple[npt.NDArray[np.float64], ...])
245assert_type(np.atleast_1d(C, C, C), tuple[npt.NDArray[Any], ...])
246
247assert_type(np.atleast_2d(A), npt.NDArray[np.float64])
248assert_type(np.atleast_2d(A, A), tuple[npt.NDArray[np.float64], npt.NDArray[np.float64]])
249assert_type(np.atleast_2d(A, A, A), tuple[npt.NDArray[np.float64], ...])
250
251assert_type(np.atleast_3d(A), npt.NDArray[np.float64])
252assert_type(np.atleast_3d(A, A), tuple[npt.NDArray[np.float64], npt.NDArray[np.float64]])
253assert_type(np.atleast_3d(A, A, A), tuple[npt.NDArray[np.float64], ...])
254
255assert_type(np.vstack([A, A]), npt.NDArray[np.float64])
256assert_type(np.vstack([A, A], dtype=np.float32), npt.NDArray[np.float32])
257assert_type(np.vstack([A, C]), npt.NDArray[Any])
258assert_type(np.vstack([C, C]), npt.NDArray[Any])
259
260assert_type(np.hstack([A, A]), npt.NDArray[np.float64])
261assert_type(np.hstack([A, A], dtype=np.float32), npt.NDArray[np.float32])
262
263assert_type(np.stack([A, A]), npt.NDArray[np.float64])
264assert_type(np.stack([A, A], dtype=np.float32), npt.NDArray[np.float32])
265assert_type(np.stack([A, C]), npt.NDArray[Any])
266assert_type(np.stack([C, C]), npt.NDArray[Any])
267assert_type(np.stack([A, A], axis=0), npt.NDArray[np.float64])
268assert_type(np.stack([A, A], out=B), SubClass[np.float64])
269
270assert_type(np.block([[A, A], [A, A]]), npt.NDArray[Any])  # pyright correctly infers this as NDArray[float64]
271assert_type(np.block(C), npt.NDArray[Any])
272
273if sys.version_info >= (3, 12):
274    from collections.abc import Buffer
275
276    def create_array(obj: npt.ArrayLike) -> npt.NDArray[Any]: ...
277
278    buffer: Buffer
279    assert_type(create_array(buffer), npt.NDArray[Any])
280