ENH: Added missing methods to `np.flatiter` (#17180)

* MAINT: Added missing methods to `np.flatiter`

* DOC: Added a comment about weird `flatiter.__getitem__` behavior

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)

3 files changed, 53 insertions, 0 deletions

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*