diff options
Diffstat (limited to 'numpy')
-rw-r--r-- | numpy/__init__.pyi | 16 | ||||
-rw-r--r-- | numpy/tests/typing/fail/flatiter.py | 25 | ||||
-rw-r--r-- | numpy/tests/typing/pass/flatiter.py | 14 | ||||
-rw-r--r-- | numpy/tests/typing/reveal/flatiter.py | 14 |
4 files changed, 69 insertions, 0 deletions
diff --git a/numpy/__init__.pyi b/numpy/__init__.pyi index d7f272616..4f7d18cdf 100644 --- a/numpy/__init__.pyi +++ b/numpy/__init__.pyi @@ -144,6 +144,14 @@ class _flagsobj: def __getitem__(self, key: str) -> bool: ... def __setitem__(self, key: str, value: bool) -> None: ... +_ArrayLikeInt = Union[ + int, + integer, + Sequence[Union[int, integer]], + Sequence[Sequence[Any]], # TODO: wait for support for recursive types + ndarray +] + _FlatIterSelf = TypeVar("_FlatIterSelf", bound=flatiter) class flatiter(Generic[_ArraySelf]): @@ -156,6 +164,14 @@ class flatiter(Generic[_ArraySelf]): def copy(self) -> _ArraySelf: ... def __iter__(self: _FlatIterSelf) -> _FlatIterSelf: ... def __next__(self) -> generic: ... + def __len__(self) -> int: ... + @overload + def __getitem__(self, key: Union[int, integer]) -> generic: ... + @overload + def __getitem__( + self, key: Union[_ArrayLikeInt, slice, ellipsis], + ) -> _ArraySelf: ... + def __array__(self, __dtype: DtypeLike = ...) -> ndarray: ... _OrderKACF = Optional[Literal["K", "A", "C", "F"]] _OrderACF = Optional[Literal["A", "C", "F"]] diff --git a/numpy/tests/typing/fail/flatiter.py b/numpy/tests/typing/fail/flatiter.py new file mode 100644 index 000000000..e8a82344f --- /dev/null +++ b/numpy/tests/typing/fail/flatiter.py @@ -0,0 +1,25 @@ +from typing import Any + +import numpy as np +from numpy.typing import DtypeLike, _SupportsArray + + +class Index: + def __index__(self) -> int: + ... + + +a: "np.flatiter[np.ndarray]" +supports_array: _SupportsArray + +a.base = Any # E: Property "base" defined in "flatiter" is read-only +a.coords = Any # E: Property "coords" defined in "flatiter" is read-only +a.index = Any # E: Property "index" defined in "flatiter" is read-only +a.copy(order='C') # E: Unexpected keyword argument + +# NOTE: Contrary to `ndarray.__getitem__` its counterpart in `flatiter` +# does not accept objects with the `__array__` or `__index__` protocols; +# boolean indexing is just plain broken (gh-17175) +a[np.bool_()] # E: No overload variant of "__getitem__" +a[Index()] # E: No overload variant of "__getitem__" +a[supports_array] # E: No overload variant of "__getitem__" diff --git a/numpy/tests/typing/pass/flatiter.py b/numpy/tests/typing/pass/flatiter.py new file mode 100644 index 000000000..93c15f601 --- /dev/null +++ b/numpy/tests/typing/pass/flatiter.py @@ -0,0 +1,14 @@ +import numpy as np + +a = np.random.rand(5).flat + +a.base +a.copy() +a.coords +a.index +iter(a) +next(a) +a[0] +a[[0, 1, 2]] +a[...] +a[:] diff --git a/numpy/tests/typing/reveal/flatiter.py b/numpy/tests/typing/reveal/flatiter.py new file mode 100644 index 000000000..56cdc7a0e --- /dev/null +++ b/numpy/tests/typing/reveal/flatiter.py @@ -0,0 +1,14 @@ +import numpy as np + +a: "np.flatiter[np.ndarray]" + +reveal_type(a.base) # E: numpy.ndarray* +reveal_type(a.copy()) # E: numpy.ndarray* +reveal_type(a.coords) # E: tuple[builtins.int] +reveal_type(a.index) # E: int +reveal_type(iter(a)) # E: Iterator[numpy.generic*] +reveal_type(next(a)) # E: numpy.generic +reveal_type(a[0]) # E: numpy.generic +reveal_type(a[[0, 1, 2]]) # E: numpy.ndarray* +reveal_type(a[...]) # E: numpy.ndarray* +reveal_type(a[:]) # E: numpy.ndarray* |