8000 Merge pull request #18545 from BvB93/arrayterator · numpy/numpy@cb7c292 · GitHub
[go: up one dir, main page]
Skip to content

Commit cb7c292

Browse files
rgommersrgommers
authored
Merge pull request #18545 from BvB93/arrayterator
ENH: Add annotations for `np.lib.arrayterator`
2 parents c5de5b5 + 025dcca commit cb7c292

File tree

12 files changed

+138
-18
lines changed

12 files changed

+138
-18
lines changed

‎numpy/__init__.pyi

Lines changed: 3 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -853,7 +853,7 @@ class flatiter(Generic[_NdArraySubClass]):
853853
@overload
854854
def __array__(self: flatiter[ndarray[Any, _DType]], __dtype: None = ...) -> ndarray[Any, _DType]: ...
855855
@overload
856-
def __array__(self, __dtype: DTypeLike) -> ndarray[Any, dtype[Any]]: ...
856+
def __array__(self, __dtype: _DType) -> ndarray[Any, _DType]: ...
857857

858858
_OrderKACF = Optional[Literal["K", "A", "C", "F"]]
859859
_OrderACF = Optional[Literal["A", "C", "F"]]
@@ -1362,7 +1362,7 @@ class ndarray(_ArrayOrScalarCommon, Generic[_ShapeType, _DType_co]):
13621362
@overload
13631363
def __array__(self, __dtype: None = ...) -> ndarray[Any, _DType_co]: ...
13641364
@overload
1365-
def __array__(self, __dtype: DTypeLike) -> ndarray[Any, dtype[Any]]: ...
1365+
def __array__(self, __dtype: _DType) -> ndarray[Any, _DType]: ...
13661366
@property
13671367
def ctypes(self) -> _ctypes: ...
13681368
@property
@@ -2409,7 +2409,7 @@ class generic(_ArrayOrScalarCommon):
24092409
@overload
24102410
def __array__(self: _ScalarType, __dtype: None = ...) -> ndarray[Any, dtype[_ScalarType]]: ...
24112411
@overload
2412-
def __array__(self, __dtype: DTypeLike) -> ndarray[Any, dtype[Any]]: ...
2412+
def __array__(self, __dtype: _DType) -> ndarray[Any, _DType]: ...
24132413
@property
24142414
def base(self) -> None: ...
24152415
@property

‎numpy/lib/__init__.pyi

Lines changed: 4 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -5,6 +5,10 @@ from numpy import (
55
ndindex as ndindex,
66
)
77

8+
from numpy.lib.arrayterator import (
9+
Arrayterator as Arrayterator,
10+
)
11+
812
from numpy.lib.index_tricks import (
913
ravel_multi_index as ravel_multi_index,
1014
unravel_index as unravel_index,
< 6D40 code class="diff-text-cell hunk">
@@ -31,7 +35,6 @@ __all__: List[str]
3135
emath: Any
3236
math: Any
3337
tracemalloc_domain: Any
34-
Arrayterator: Any
3538
iscomplexobj: Any
3639
isrealobj: Any
3740
imag: Any

‎numpy/lib/arrayterator.pyi

Lines changed: 53 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,53 @@
1+
import sys
2+
from typing import (
3+
List,
4+
Any,
5+
TypeVar,
6+
Generator,
7+
List,
8+
Union,
9+
Tuple,
10+
overload,
11+
)
12+
13+
from numpy import ndarray, dtype, generic
14+
from numpy.typing import DTypeLike
15+
16+
# TODO: Set a shape bound once we've got proper shape support
17+
_Shape = TypeVar("_Shape", bound=Any)
18+
_DType = TypeVar("_DType", bound=dtype[Any])
19+
_ScalarType = TypeVar("_ScalarType", bound=generic)
20+
21+
_Index = Union[
22+
Union[ellipsis, int, slice],
23+
Tuple[Union[ellipsis, int, slice], ...],
24+
]
25+
26+
__all__: List[str]
27+
28+
# NOTE: In reality `Arrayterator` does not actually inherit from `ndarray`,
29+
# but its ``__getattr__` method does wrap around the former and thus has
30+
# access to all its methods
31+
32+
class Arrayterator(ndarray[_Shape, _DType]):
33+
var: ndarray[_Shape, _DType] # type: ignore[assignment]
34+
buf_size: None | int
35+
start: List[int]
36+
stop: List[int]
37+
step: List[int]
38+
39+
@property # type: ignore[misc]
40+
def shape(self) -> Tuple[int, ...]: ...
41+
@property
42+
def flat( # type: ignore[override]
43+
self: ndarray[Any, dtype[_ScalarType]]
44+
) -> Generator[_ScalarType, None, None]: ...
45+
def __init__(
46+
self, var: ndarray[_Shape, _DType], buf_size: None | int = ...
47+
) -> None: ...
48+
@overload
49+
def __array__(self, dtype: None = ...) -> ndarray[Any, _DType]: ...
50+
@overload
51+
def __array__(self, dtype: DTypeLike) -> ndarray[Any, dtype[Any]]: ...
52+
def __getitem__(self, index: _Index) -> Arrayterator[Any, _DType]: ...
53+
def __iter__(self) -> Generator[ndarray[Any, _DType], None, None]: ...

‎numpy/typing/_array_like.py

Lines changed: 3 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -40,11 +40,12 @@
4040

4141
if TYPE_CHECKING or HAVE_PROTOCOL:
4242
# The `_SupportsArray` protocol only cares about the default dtype
43-
# (i.e. `dtype=None`) of the to-be returned array.
43+
# (i.e. `dtype=None` or no `dtype` parameter at all) of the to-be returned
44+
# array.
4445
# Concrete implementations of the protocol are responsible for adding
4546
# any and all remaining overloads
4647
class _SupportsArray(Protocol[_DType_co]):
47-
def __array__(self, dtype: None = ...) -> ndarray[Any, _DType_co]: ...
48+
def __array__(self) -> ndarray[Any, _DType_co]: ...
4849
else:
4950
_SupportsArray = Any
5051

‎numpy/typing/tests/data/fail/arrayterator.py

Lines changed: 14 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,14 @@
1+
from typing import Any
2+
import numpy as np
3+
4+
AR_i8: np.ndarray[Any, np.dtype[np.int64]]
5+
ar_iter = np.lib.Arrayterator(AR_i8)
6+
7+
np.lib.Arrayterator(np.int64()) # E: incompatible type
8+
ar_iter.shape = (10, 5) # E: is read-only
9+
ar_iter[None] # E: Invalid index type
10+
ar_iter[None, 1] # E: Invalid index type
11+
ar_iter[np.intp()] # E: Invalid index type
12+
ar_iter[np.intp(), ...] # E: Invalid index type
13+
ar_iter[AR_i8] # E: Invalid index type
14+
ar_iter[AR_i8, :] # E: Invalid index type

‎numpy/typing/tests/data/pass/arithmetic.py

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -25,7 +25,7 @@
2525

2626

2727
class Object:
28-
def __array__(self, dtype: None = None) -> np.ndarray[Any, np.dtype[np.object_]]:
28+
def __array__(self) -> np.ndarray[Any, np.dtype[np.object_]]:
2929
ret = np.empty((), dtype=object)
3030
ret[()] = self
3131
return ret

‎numpy/typing/tests/data/pass/array_like.py

Lines changed: 6 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -1,7 +1,7 @@
11
from typing import Any, List, Optional
22

33
import numpy as np
4-
from numpy.typing import ArrayLike, DTypeLike, _SupportsArray
4+
from numpy.typing import ArrayLike, _SupportsArray
55

66
x1: ArrayLike = True
77
x2: ArrayLike = 5
@@ -18,20 +18,20 @@
1818

1919

2020
class A:
21-
def __array__(self, dtype: DTypeLike = None) -> np.ndarray:
21+
def __array__(self, dtype: Optional[np.dtype] = None) -> np.ndarray:
2222
return np.array([1, 2, 3])
2323

2424

2525
x13: ArrayLike = A()
2626

2727
scalar: _SupportsArray = np.int64(1)
28-
scalar.__array__(None)
28+
scalar.__array__()
2929
array: _SupportsArray = np.array(1)
30-
array.__array__(None)
30+
array.__array__()
3131

3232
a: _SupportsArray = A()
33-
a.__array__(None)
34-
a.__array__(dtype=None)
33+
a.__array__()
34+
a.__array__()
3535

3636
# Escape hatch for when you mean to make something like an object
3737
# array.

‎numpy/typing/tests/data/pass/arrayterator.py

Lines changed: 27 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,27 @@
1+
2+
from __future__ import annotations
3+
4+
from typing import Any
5+
import numpy as np
6+
7+
AR_i8: np.ndarray[Any, np.dtype[np.int_]] = np.arange(10)
8+
ar_iter = np.lib.Arrayterator(AR_i8)
9+
10+
ar_iter.var
11+
ar_iter.buf_size
12+
ar_iter.start
13+
ar_iter.stop
14+
ar_iter.step
15+
ar_iter.shape
16+
ar_iter.flat
17+
18+
ar_iter.__array__()
19+
20+
for i in ar_iter:
21+
pass
22+
23+
ar_iter[0]
24+
ar_iter[...]
25+
ar_iter[:]
26+
ar_iter[0, 0, 0]
27+
ar_iter[..., 0, :]

‎numpy/typing/tests/data/pass/flatiter.py

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -13,4 +13,4 @@
1313
a[...]
1414
a[:]
1515
a.__array__()
16-
a.__array__(np.float64)
16+
a.__array__(np.dtype(np.float64))

‎numpy/typing/tests/data/pass/ufunclike.py

Lines changed: 1 addition & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -13,9 +13,7 @@ def __floor__(self) -> Object:
1313
def __ge__(self, value: object) -> bool:
1414
return True
1515

16-
def __array__(
17-
self, dtype: None = None
18-
) -> np.ndarray[Any, np.dtype[np.object_]]:
16+
def __array__(self) -> np.ndarray[Any, np.dtype[np.object_]]:
1917
ret = np.empty((), dtype=object)
2018
ret[()] = self
2119
return ret

‎numpy/typing/tests/data/reveal/arrayterator.py

Lines changed: 24 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,24 @@
1+
from typing import Any
2+
import numpy as np
3+
4+
AR_i8: np.ndarray[Any, np.dtype[np.int64]]
5+
ar_iter = np.lib.Arrayterator(AR_i8)
6+
7+
reveal_type(ar_iter.var) # E: numpy.ndarray[Any, numpy.dtype[{int64}]]
8+
reveal_type(ar_iter.buf_size) # E: Union[None, builtins.int]
9+
reveal_type(ar_iter.start) # E: builtins.list[builtins.int]
10+
reveal_type(ar_iter.stop) # E: builtins.list[builtins.int]
11+
reveal_type(ar_iter.step) # E: builtins.list[builtins.int]
12+
reveal_type(ar_iter.shape) # E: builtins.tuple[builtins.int]
13+
reveal_type(ar_iter.flat) # E: 'typing.Generator[{int64}, None, None]
14+
15+
reveal_type(ar_iter.__array__()) # E: numpy.ndarray[Any, numpy.dtype[{int64}]]
16+
17+
for i in ar_iter:
18+
reveal_type(i) # E: numpy.ndarray[Any, numpy.dtype[{int64}]]
19+
20+
reveal_type(ar_iter[0]) # E: numpy.lib.arrayterator.Arrayterator[Any, numpy.dtype[{int64}]]
21+
reveal_type(ar_iter[...]) # E: numpy.lib.arrayterator.Arrayterator[Any, numpy.dtype[{int64}]]
22+
reveal_type(ar_iter[:]) # E: numpy.lib.arrayterator.Arrayterator[Any, numpy.dtype[{int64}]]
23+
reveal_type(ar_iter[0, 0, 0]) # E: numpy.lib.arrayterator.Arrayterator[Any, numpy.dtype[{int64}]]
24+
reveal_type(ar_iter[..., 0, :]) # E: numpy.lib.arrayterator.Arrayterator[Any, numpy.dtype[{int64}]]

‎numpy/typing/tests/data/reveal/flatiter.py

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -14,4 +14,4 @@
1414
reveal_type(a[...]) # E: numpy.ndarray[Any, numpy.dtype[numpy.str_]]
1515
reveal_type(a[:]) # E: numpy.ndarray[Any, numpy.dtype[numpy.str_]]
1616
reveal_type(a.__array__()) # E: numpy.ndarray[Any, numpy.dtype[numpy.str_]]
17-
reveal_type(a.__array__(np.float64)) # E: numpy.ndarray[Any, numpy.dtype[Any]]
17+
reveal_type(a.__array__(np.dtype(np.float64))) # E: numpy.ndarray[Any, numpy.dtype[{float64}]]

0 commit comments

Comments
 (0)
0