1
1
# TODO: Sort out any and all missing functions in this namespace
2
- import builtins
3
2
import os
4
3
import datetime as dt
5
4
from collections .abc import Sequence , Callable , Iterable
@@ -8,16 +7,20 @@ from typing import (
8
7
Any ,
9
8
TypeAlias ,
10
9
overload ,
10
+ TypeAlias ,
11
11
TypeVar ,
12
+ TypedDict ,
12
13
SupportsIndex ,
13
14
final ,
14
15
Final ,
15
16
Protocol ,
16
17
ClassVar ,
18
+ type_check_only ,
17
19
)
20
+ from typing_extensions import Unpack
18
21
19
22
import numpy as np
20
- from numpy import (
23
+ from numpy import ( # type: ignore[attr-defined]
21
24
# Re-exports
22
25
busdaycalendar as busdaycalendar ,
23
26
broadcast as broadcast ,
@@ -57,6 +60,7 @@ from numpy._typing import (
57
60
# DTypes
58
61
DTypeLike ,
59
62
_DTypeLike ,
63
+ _SupportsDType ,
60
64
61
65
# Arrays
62
66
NDArray ,
@@ -90,6 +94,7 @@ from numpy._typing._ufunc import (
90
94
_T_co = TypeVar ("_T_co" , covariant = True )
91
95
_T_contra = TypeVar ("_T_contra" , contravariant = True )
92
96
_SCT = TypeVar ("_SCT" , bound = generic )
97
+ _DType = TypeVar ("_DType" , bound = np .dtype [Any ])
93
98
_ArrayType = TypeVar ("_ArrayType" , bound = ndarray [Any , Any ])
94
99
_ArrayType_co = TypeVar (
95
100
"_ArrayType_co" ,
@@ -102,7 +107,9 @@ _Nin = TypeVar("_Nin", bound=int)
102
107
_Nout = TypeVar ("_Nout" , bound = int )
103
108
104
109
_SizeType = TypeVar ("_SizeType" , bound = int )
110
+ _ShapeType = TypeVar ("_ShapeType" , bound = tuple [int , ...])
105
111
_1DArray : TypeAlias = ndarray [tuple [_SizeType ], dtype [_SCT ]]
112
+ _Array : TypeAlias = ndarray [_ShapeType , dtype [_SCT ]]
106
113
107
114
# Valid time units
108
115
_UnitKind = L [
@@ -136,6 +143,119 @@ class _SupportsLenAndGetItem(Protocol[_T_contra, _T_co]):
136
143
class _SupportsArray (Protocol [_ArrayType_co ]):
137
144
def __array__ (self , / ) -> _ArrayType_co : ...
138
145
146
+ @type_check_only
147
+ class _KwargsEmptyLike (TypedDict , total = False ):
148
+ device : None | L ["cpu" ]
149
+
150
+ @type_check_only
151
+ class _KwargsEmpty (_KwargsEmptyLike , total = False ):
152
+ like : None | _SupportsArrayFunc
153
+
154
+ @type_check_only
155
+ class _ConstructorEmpty (Protocol ):
156
+ # 1-D shape
157
+ @overload
158
+ def __call__ (
159
+ self , / ,
160
+ shape : _SizeType ,
161
+ dtype : None = ...,
162
+ order : _OrderCF = ...,
163
+ ** kwargs : Unpack [_KwargsEmpty ],
164
+ ) -> _Array [tuple [_SizeType ], float64 ]: ...
165
+ @overload
166
+ def __call__ (
167
+ self , / ,
168
+ shape : _SizeType ,
169
+ dtype : _DType | _SupportsDType [_DType ],
170
+ order : _OrderCF = ...,
171
+ ** kwargs : Unpack [_KwargsEmpty ],
172
+ ) -> ndarray [tuple [_SizeType ], _DType ]: ...
173
+ @overload
174
+ def __call__ (
175
+ self , / ,
176
+ shape : _SizeType ,
177
+ dtype : type [_SCT ],
178
+ order : _OrderCF = ...,
179
+ ** kwargs : Unpack [_KwargsEmpty ],
180
+ ) -> _Array [tuple [_SizeType ], _SCT ]: ...
181
+ @overload
182
+ def __call__ (
183
+ self , / ,
184
+ shape : _SizeType ,
185
+ dtype : DTypeLike ,
186
+ order : _OrderCF = ...,
187
+ ** kwargs : Unpack [_KwargsEmpty ],
188
+ ) -> _Array [tuple [_SizeType ], Any ]: ...
189
+
190
+ # known shape
191
+ @overload
192
+ def __call__ (
193
+ self , / ,
194
+ shape : _ShapeType ,
195
+ dtype : None = ...,
196
+ order : _OrderCF = ...,
197
+ ** kwargs : Unpack [_KwargsEmpty ],
198
+ ) -> _Array [_ShapeType , float64 ]: ...
199
+ @overload
200
+ def __call__ (
201
+ self , / ,
202
+ shape : _ShapeType ,
203
+ dtype : _DType | _SupportsDType [_DType ],
204
+ order : _OrderCF = ...,
205
+ ** kwargs : Unpack [_KwargsEmpty ],
206
+ ) -> ndarray [_ShapeType , _DType ]: ...
207
+ @overload
208
+ def __call__ (
209
+ self , / ,
210
+ shape : _ShapeType ,
211
+ dtype : type [_SCT ],
212
+ order : _OrderCF = ...,
213
+ ** kwargs : Unpack [_KwargsEmpty ],
214
+ ) -> _Array [_ShapeType , _SCT ]: ...
215
+ @overload
216
+ def __call__ (
217
+ self , / ,
218
+ shape : _ShapeType ,
219
+ dtype : DTypeLike ,
220
+ order : _OrderCF = ...,
221
+ ** kwargs : Unpack [_KwargsEmpty ],
222
+ ) -> _Array [_ShapeType , Any ]: ...
223
+
224
+ # unknown shape
225
+ @overload
226
+ def __call__ (
227
+ self , / ,
228
+ shape : _ShapeLike ,
229
+ dtype : None = ...,
230
+ order : _OrderCF = ...,
231
+ ** kwargs : Unpack [_KwargsEmpty ],
232
+ ) -> NDArray [float64 ]: ...
233
+ @overload
234
+ def __call__ (
235
+ self , / ,
236
+ shape : _ShapeLike ,
237
+ dtype : _DType | _SupportsDType [_DType ],
238
+ order : _OrderCF = ...,
239
+ ** kwargs : Unpack [_KwargsEmpty ],
240
+ ) -> ndarray [Any , _DType ]: ...
241
+ @overload
242
+ def __call__ (
243
+ self , / ,
244
+ shape : _ShapeLike ,
245
+ dtype : type [_SCT ],
246
+ order : _OrderCF = ...,
247
+ ** kwargs : Unpack [_KwargsEmpty ],
248
+ ) -> NDArray [_SCT ]: ...
249
+ @overload
250
+ def __call__ (
251
+ self , / ,
252
+ shape : _ShapeLike ,
253
+ dtype : DTypeLike ,
254
+ order : _OrderCF = ...,
255
+ ** kwargs : Unpack [_KwargsEmpty ],
256
+ ) -> NDArray [Any ]: ...
257
+
258
+
139
259
__all__ : list [str ]
140
260
141
261
ALLOW_THREADS : Final [int ] # 0 or 1 (system-specific)
@@ -148,6 +268,9 @@ MAY_SHARE_BOUNDS: L[0]
148
268
MAY_SHARE_EXACT : L [- 1 ]
149
269
tracemalloc_domain : L [389047 ]
150
270
271
+ zeros : Final [_ConstructorEmpty ]
272
+ empty : Final [_ConstructorEmpty ]
273
+
151
274
@overload
152
275
def empty_like (
153
276
prototype : _ArrayType ,
@@ -266,62 +389,6 @@ def array(
266
389
like : None | _SupportsArrayFunc = ...,
267
390
) -> NDArray [Any ]: ...
268
391
269
- @overload
270
- def zeros (
271
- shape : _ShapeLike ,
272
- dtype : None = ...,
273
- order : _OrderCF = ...,
274
- * ,
275
- device : None | L ["cpu" ] = ...,
276
- like : None | _SupportsArrayFunc = ...,
277
- ) -> NDArray [float64 ]: ...
278
- @overload
279
- def zeros (
280
- shape : _ShapeLike ,
281
- dtype : _DTypeLike [_SCT ],
282
- order : _OrderCF = ...,
283
- * ,
284
- device : None | L ["cpu" ] = ...,
285
- like : None | _SupportsArrayFunc = ...,
286
- ) -> NDArray [_SCT ]: ...
287
- @overload
288
- def zeros (
289
- shape : _ShapeLike ,
290
- dtype : DTypeLike ,
291
- order : _OrderCF = ...,
292
- * ,
293
- device : None | L ["cpu" ] = ...,
294
- like : None | _SupportsArrayFunc = ...,
295
- ) -> NDArray [Any ]: ...
296
-
297
- @overload
298
- def empty (
299
- shape : _ShapeLike ,
300
- dtype : None = ...,
301
- order : _OrderCF = ...,
302
- * ,
303
- device : None | L ["cpu" ] = ...,
304
- like : None | _SupportsArrayFunc = ...,
305
- ) -> NDArray [float64 ]: ...
306
- @overload
307
- def empty (
308
- shape : _ShapeLike ,
309
- dtype : _DTypeLike [_SCT ],
310
- order : _OrderCF = ...,
311
- * ,
312
- device : None | L ["cpu" ] = ...,
313
- like : None | _SupportsArrayFunc = ...,
314
- ) -> NDArray [_SCT ]: ...
315
- @overload
316
- def empty (
317
- shape : _ShapeLike ,
318
- dtype : DTypeLike ,
319
- order : _OrderCF = ...,
320
- * ,
321
- device : None | L ["cpu" ] = ...,
322
- like : None | _SupportsArrayFunc = ...,
323
- ) -> NDArray [Any ]: ...
324
-
325
392
@overload
326
393
def unravel_index ( # type: ignore[misc]
327
394
indices : _IntLike_co ,
0 commit comments