codekingpro/portable-devtools

1from typing import Any, TypeAlias, TypeVar, assert_type, type_check_only
2
3import numpy as np
4import numpy.typing as npt
5
6_ScalarT = TypeVar("_ScalarT", bound=np.generic)
7
8_1D: TypeAlias = tuple[int]
9_2D: TypeAlias = tuple[int, int]
10_ND: TypeAlias = tuple[Any, ...]
11
12_Indices2D: TypeAlias = tuple[
13    np.ndarray[_1D, np.dtype[np.intp]],
14    np.ndarray[_1D, np.dtype[np.intp]],
15]
16
17###
18
19_nd_bool: np.ndarray[_ND, np.dtype[np.bool]]
20_1d_bool: np.ndarray[_1D, np.dtype[np.bool]]
21_2d_bool: np.ndarray[_2D, np.dtype[np.bool]]
22_nd_u64: np.ndarray[_ND, np.dtype[np.uint64]]
23_nd_i64: np.ndarray[_ND, np.dtype[np.int64]]
24_nd_f64: np.ndarray[_ND, np.dtype[np.float64]]
25_nd_c128: np.ndarray[_ND, np.dtype[np.complex128]]
26_nd_obj: np.ndarray[_ND, np.dtype[np.object_]]
27
28_to_nd_bool: list[bool] | list[list[bool]]
29_to_1d_bool: list[bool]
30_to_2d_bool: list[list[bool]]
31
32_to_1d_f64: list[float]
33_to_1d_c128: list[complex]
34
35@type_check_only
36def func1(ar: npt.NDArray[_ScalarT], a: int) -> npt.NDArray[_ScalarT]: ...
37@type_check_only
38def func2(ar: npt.NDArray[np.number], a: str) -> npt.NDArray[np.float64]: ...
39
40@type_check_only
41class _Cube:
42    shape = 3, 4
43    ndim = 2
44
45###
46
47# fliplr
48assert_type(np.fliplr(_nd_bool), np.ndarray[_ND, np.dtype[np.bool]])
49assert_type(np.fliplr(_1d_bool), np.ndarray[_1D, np.dtype[np.bool]])
50assert_type(np.fliplr(_2d_bool), np.ndarray[_2D, np.dtype[np.bool]])
51assert_type(np.fliplr(_to_nd_bool), np.ndarray)
52assert_type(np.fliplr(_to_1d_bool), np.ndarray)
53assert_type(np.fliplr(_to_2d_bool), np.ndarray)
54
55# flipud
56assert_type(np.flipud(_nd_bool), np.ndarray[_ND, np.dtype[np.bool]])
57assert_type(np.flipud(_1d_bool), np.ndarray[_1D, np.dtype[np.bool]])
58assert_type(np.flipud(_2d_bool), np.ndarray[_2D, np.dtype[np.bool]])
59assert_type(np.flipud(_to_nd_bool), np.ndarray)
60assert_type(np.flipud(_to_1d_bool), np.ndarray)
61assert_type(np.flipud(_to_2d_bool), np.ndarray)
62
63# eye
64assert_type(np.eye(10), np.ndarray[_2D, np.dtype[np.float64]])
65assert_type(np.eye(10, M=20, dtype=np.int64), np.ndarray[_2D, np.dtype[np.int64]])
66assert_type(np.eye(10, k=2, dtype=int), np.ndarray[_2D])
67
68# diag
69assert_type(np.diag(_nd_bool), np.ndarray[_ND, np.dtype[np.bool]])
70assert_type(np.diag(_1d_bool), np.ndarray[_2D, np.dtype[np.bool]])
71assert_type(np.diag(_2d_bool), np.ndarray[_1D, np.dtype[np.bool]])
72assert_type(np.diag(_to_nd_bool, k=0), np.ndarray)
73assert_type(np.diag(_to_1d_bool, k=0), np.ndarray[_2D])
74assert_type(np.diag(_to_2d_bool, k=0), np.ndarray[_1D])
75
76# diagflat
77assert_type(np.diagflat(_nd_bool), np.ndarray[_2D, np.dtype[np.bool]])
78assert_type(np.diagflat(_1d_bool), np.ndarray[_2D, np.dtype[np.bool]])
79assert_type(np.diagflat(_2d_bool), np.ndarray[_2D, np.dtype[np.bool]])
80assert_type(np.diagflat(_to_nd_bool, k=0), np.ndarray[_2D])
81assert_type(np.diagflat(_to_1d_bool, k=0), np.ndarray[_2D])
82assert_type(np.diagflat(_to_2d_bool, k=0), np.ndarray[_2D])
83
84# tri
85assert_type(np.tri(10), np.ndarray[_2D, np.dtype[np.float64]])
86assert_type(np.tri(10, M=20, dtype=np.int64), np.ndarray[_2D, np.dtype[np.int64]])
87assert_type(np.tri(10, k=2, dtype=int), np.ndarray[_2D])
88
89# tril
90assert_type(np.tril(_nd_bool), np.ndarray[_ND, np.dtype[np.bool]])
91assert_type(np.tril(_to_nd_bool, k=0), np.ndarray)
92assert_type(np.tril(_to_1d_bool, k=0), np.ndarray)
93assert_type(np.tril(_to_2d_bool, k=0), np.ndarray)
94
95# triu
96assert_type(np.triu(_nd_bool), np.ndarray[_ND, np.dtype[np.bool]])
97assert_type(np.triu(_to_nd_bool, k=0), np.ndarray)
98assert_type(np.triu(_to_1d_bool, k=0), np.ndarray)
99assert_type(np.triu(_to_2d_bool, k=0), np.ndarray)
100
101# vander
102assert_type(np.vander(_nd_bool), np.ndarray[_2D, np.dtype[np.int_]])
103assert_type(np.vander(_nd_u64), np.ndarray[_2D, np.dtype[np.uint64]])
104assert_type(np.vander(_nd_i64, N=2), np.ndarray[_2D, np.dtype[np.int64]])
105assert_type(np.vander(_nd_f64, increasing=True), np.ndarray[_2D, np.dtype[np.float64]])
106assert_type(np.vander(_nd_c128), np.ndarray[_2D, np.dtype[np.complex128]])
107assert_type(np.vander(_nd_obj), np.ndarray[_2D, np.dtype[np.object_]])
108
109# histogram2d
110assert_type(
111    np.histogram2d(_to_1d_f64, _to_1d_f64),
112    tuple[
113        np.ndarray[_2D, np.dtype[np.float64]],
114        np.ndarray[_1D, np.dtype[np.float64]],
115        np.ndarray[_1D, np.dtype[np.float64]],
116    ],
117)
118assert_type(
119    np.histogram2d(_to_1d_c128, _to_1d_c128),
120    tuple[
121        np.ndarray[_2D, np.dtype[np.float64]],
122        np.ndarray[_1D, np.dtype[np.complex128 | Any]],
123        np.ndarray[_1D, np.dtype[np.complex128 | Any]],
124    ],
125)
126assert_type(
127    np.histogram2d(_nd_i64, _nd_bool),
128    tuple[
129        np.ndarray[_2D, np.dtype[np.float64]],
130        np.ndarray[_1D, np.dtype[np.float64]],
131        np.ndarray[_1D, np.dtype[np.float64]],
132    ],
133)
134assert_type(
135    np.histogram2d(_nd_f64, _nd_i64),
136    tuple[
137        np.ndarray[_2D, np.dtype[np.float64]],
138        np.ndarray[_1D, np.dtype[np.float64]],
139        np.ndarray[_1D, np.dtype[np.float64]],
140    ],
141)
142assert_type(
143    np.histogram2d(_nd_i64, _nd_f64),
144    tuple[
145        np.ndarray[_2D, np.dtype[np.float64]],
146        np.ndarray[_1D, np.dtype[np.float64]],
147        np.ndarray[_1D, np.dtype[np.float64]],
148    ],
149)
150assert_type(
151    np.histogram2d(_nd_f64, _nd_c128, weights=_to_1d_bool),
152    tuple[
153        np.ndarray[_2D, np.dtype[np.float64]],
154        np.ndarray[_1D, np.dtype[np.complex128]],
155        np.ndarray[_1D, np.dtype[np.complex128]],
156    ],
157)
158assert_type(
159    np.histogram2d(_nd_f64, _nd_c128, bins=8),
160    tuple[
161        np.ndarray[_2D, np.dtype[np.float64]],
162        np.ndarray[_1D, np.dtype[np.complex128]],
163        np.ndarray[_1D, np.dtype[np.complex128]],
164    ],
165)
166assert_type(
167    np.histogram2d(_nd_c128, _nd_f64, bins=(8, 5)),
168    tuple[
169        np.ndarray[_2D, np.dtype[np.float64]],
170        np.ndarray[_1D, np.dtype[np.complex128]],
171        np.ndarray[_1D, np.dtype[np.complex128]],
172    ],
173)
174assert_type(
175    np.histogram2d(_nd_c128, _nd_i64, bins=_nd_u64),
176    tuple[
177        np.ndarray[_2D, np.dtype[np.float64]],
178        np.ndarray[_1D, np.dtype[np.uint64]],
179        np.ndarray[_1D, np.dtype[np.uint64]],
180    ],
181)
182assert_type(
183    np.histogram2d(_nd_c128, _nd_c128, bins=(_nd_u64, _nd_u64)),
184    tuple[
185        np.ndarray[_2D, np.dtype[np.float64]],
186        np.ndarray[_1D, np.dtype[np.uint64]],
187        np.ndarray[_1D, np.dtype[np.uint64]],
188    ],
189)
190assert_type(
191    np.histogram2d(_nd_c128, _nd_c128, bins=(_nd_bool, 8)),
192    tuple[
193        np.ndarray[_2D, np.dtype[np.float64]],
194        np.ndarray[_1D, np.dtype[np.complex128 | np.bool]],
195        np.ndarray[_1D, np.dtype[np.complex128 | np.bool]],
196    ],
197)
198assert_type(
199    np.histogram2d(_nd_c128, _nd_c128, bins=(_to_1d_f64, 8)),
200    tuple[
201        np.ndarray[_2D, np.dtype[np.float64]],
202        np.ndarray[_1D, np.dtype[np.complex128 | Any]],
203        np.ndarray[_1D, np.dtype[np.complex128 | Any]],
204    ],
205)
206
207# mask_indices
208assert_type(np.mask_indices(10, func1), _Indices2D)
209assert_type(np.mask_indices(8, func2, "0"), _Indices2D)
210
211# tril_indices
212assert_type(np.tril_indices(3), _Indices2D)
213assert_type(np.tril_indices(3, 1), _Indices2D)
214assert_type(np.tril_indices(3, 1, 2), _Indices2D)
215# tril_indices
216assert_type(np.triu_indices(3), _Indices2D)
217assert_type(np.triu_indices(3, 1), _Indices2D)
218assert_type(np.triu_indices(3, 1, 2), _Indices2D)
219
220# tril_indices_from
221assert_type(np.tril_indices_from(_2d_bool), _Indices2D)
222assert_type(np.tril_indices_from(_Cube()), _Indices2D)
223# triu_indices_from
224assert_type(np.triu_indices_from(_2d_bool), _Indices2D)
225assert_type(np.triu_indices_from(_Cube()), _Indices2D)
226