Skip to content

Commit f1d11b3

Browse files
authored
Merge pull request numpy#27171 from jorenham/typing/shaped-array-constructors
TYP: Shape-typed array constructors: ``numpy.{empty,zeros,ones,full}``
2 parents b599264 + f3dd2e1 commit f1d11b3

File tree

4 files changed

+250
-112
lines changed

4 files changed

+250
-112
lines changed

numpy/_core/multiarray.pyi

Lines changed: 125 additions & 58 deletions
Original file line numberDiff line numberDiff line change
@@ -1,5 +1,4 @@
11
# TODO: Sort out any and all missing functions in this namespace
2-
import builtins
32
import os
43
import datetime as dt
54
from collections.abc import Sequence, Callable, Iterable
@@ -8,16 +7,20 @@ from typing import (
87
Any,
98
TypeAlias,
109
overload,
10+
TypeAlias,
1111
TypeVar,
12+
TypedDict,
1213
SupportsIndex,
1314
final,
1415
Final,
1516
Protocol,
1617
ClassVar,
18+
type_check_only,
1719
)
20+
from typing_extensions import Unpack
1821

1922
import numpy as np
20-
from numpy import (
23+
from numpy import ( # type: ignore[attr-defined]
2124
# Re-exports
2225
busdaycalendar as busdaycalendar,
2326
broadcast as broadcast,
@@ -57,6 +60,7 @@ from numpy._typing import (
5760
# DTypes
5861
DTypeLike,
5962
_DTypeLike,
63+
_SupportsDType,
6064

6165
# Arrays
6266
NDArray,
@@ -90,6 +94,7 @@ from numpy._typing._ufunc import (
9094
_T_co = TypeVar("_T_co", covariant=True)
9195
_T_contra = TypeVar("_T_contra", contravariant=True)
9296
_SCT = TypeVar("_SCT", bound=generic)
97+
_DType = TypeVar("_DType", bound=np.dtype[Any])
9398
_ArrayType = TypeVar("_ArrayType", bound=ndarray[Any, Any])
9499
_ArrayType_co = TypeVar(
95100
"_ArrayType_co",
@@ -102,7 +107,9 @@ _Nin = TypeVar("_Nin", bound=int)
102107
_Nout = TypeVar("_Nout", bound=int)
103108

104109
_SizeType = TypeVar("_SizeType", bound=int)
110+
_ShapeType = TypeVar("_ShapeType", bound=tuple[int, ...])
105111
_1DArray: TypeAlias = ndarray[tuple[_SizeType], dtype[_SCT]]
112+
_Array: TypeAlias = ndarray[_ShapeType, dtype[_SCT]]
106113

107114
# Valid time units
108115
_UnitKind = L[
@@ -136,6 +143,119 @@ class _SupportsLenAndGetItem(Protocol[_T_contra, _T_co]):
136143
class _SupportsArray(Protocol[_ArrayType_co]):
137144
def __array__(self, /) -> _ArrayType_co: ...
138145

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+
139259
__all__: list[str]
140260

141261
ALLOW_THREADS: Final[int] # 0 or 1 (system-specific)
@@ -148,6 +268,9 @@ MAY_SHARE_BOUNDS: L[0]
148268
MAY_SHARE_EXACT: L[-1]
149269
tracemalloc_domain: L[389047]
150270

271+
zeros: Final[_ConstructorEmpty]
272+
empty: Final[_ConstructorEmpty]
273+
151274
@overload
152275
def empty_like(
153276
prototype: _ArrayType,
@@ -266,62 +389,6 @@ def array(
266389
like: None | _SupportsArrayFunc = ...,
267390
) -> NDArray[Any]: ...
268391

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-
325392
@overload
326393
def unravel_index( # type: ignore[misc]
327394
indices: _IntLike_co,

0 commit comments

Comments
 (0)