summaryrefslogtreecommitdiff
path: root/numpy
diff options
context:
space:
mode:
Diffstat (limited to 'numpy')
-rw-r--r--numpy/__init__.pyi16
-rw-r--r--numpy/tests/typing/fail/flatiter.py25
-rw-r--r--numpy/tests/typing/pass/flatiter.py14
-rw-r--r--numpy/tests/typing/reveal/flatiter.py14
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*