TST: Add typing tests for `np.lib.arrayterator`

3 files changed, 65 insertions, 0 deletions

diff --git a/numpy/typing/tests/data/fail/arrayterator.py b/numpy/typing/tests/data/fail/arrayterator.py
new file mode 100644
index 000000000..c50fb2ec4
--- /dev/null
+++ b/numpy/typing/tests/data/fail/arrayterator.py
@@ -0,0 +1,14 @@
+from typing import Any
+import numpy as np
+
+AR_i8: np.ndarray[Any, np.dtype[np.int64]]
+ar_iter = np.lib.Arrayterator(AR_i8)
+
+np.lib.Arrayterator(np.int64())  # E: incompatible type
+ar_iter.shape = (10, 5)  # E: is read-only
+ar_iter[None]  # E: Invalid index type
+ar_iter[None, 1]  # E: Invalid index type
+ar_iter[np.intp()]  # E: Invalid index type
+ar_iter[np.intp(), ...]  # E: Invalid index type
+ar_iter[AR_i8]  # E: Invalid index type
+ar_iter[AR_i8, :]  # E: Invalid index type
diff --git a/numpy/typing/tests/data/pass/arrayterator.py b/numpy/typing/tests/data/pass/arrayterator.py
new file mode 100644
index 000000000..572be5e2f
--- /dev/null
+++ b/numpy/typing/tests/data/pass/arrayterator.py
@@ -0,0 +1,27 @@
+
+from __future__ import annotations
+
+from typing import Any
+import numpy as np
+
+AR_i8: np.ndarray[Any, np.dtype[np.int_]] = np.arange(10)
+ar_iter = np.lib.Arrayterator(AR_i8)
+
+ar_iter.var
+ar_iter.buf_size
+ar_iter.start
+ar_iter.stop
+ar_iter.step
+ar_iter.shape
+ar_iter.flat
+
+ar_iter.__array__()
+
+for i in ar_iter:
+    pass
+
+ar_iter[0]
+ar_iter[...]
+ar_iter[:]
+ar_iter[0, 0, 0]
+ar_iter[..., 0, :]
diff --git a/numpy/typing/tests/data/reveal/arrayterator.py b/numpy/typing/tests/data/reveal/arrayterator.py
new file mode 100644
index 000000000..b57861d00
--- /dev/null
+++ b/numpy/typing/tests/data/reveal/arrayterator.py
@@ -0,0 +1,24 @@
+from typing import Any
+import numpy as np
+
+AR_i8: np.ndarray[Any, np.dtype[np.int64]]
+ar_iter = np.lib.Arrayterator(AR_i8)
+
+reveal_type(ar_iter.var)  # E: numpy.ndarray[Any, numpy.dtype[{int64}]]
+reveal_type(ar_iter.buf_size)  # E: Union[None, builtins.int]
+reveal_type(ar_iter.start)  # E: builtins.list[builtins.int]
+reveal_type(ar_iter.stop)  # E: builtins.list[builtins.int]
+reveal_type(ar_iter.step)  # E: builtins.list[builtins.int]
+reveal_type(ar_iter.shape)  # E: builtins.tuple[builtins.int]
+reveal_type(ar_iter.flat)  # E: 'typing.Generator[{int64}, None, None]
+
+reveal_type(ar_iter.__array__())  # E: numpy.ndarray[Any, numpy.dtype[{int64}]]
+
+for i in ar_iter:
+    reveal_type(i)  # E: numpy.ndarray[Any, numpy.dtype[{int64}]]
+
+reveal_type(ar_iter[0])  # E: numpy.lib.arrayterator.Arrayterator[Any, numpy.dtype[{int64}]]
+reveal_type(ar_iter[...])  # E: numpy.lib.arrayterator.Arrayterator[Any, numpy.dtype[{int64}]]
+reveal_type(ar_iter[:])  # E: numpy.lib.arrayterator.Arrayterator[Any, numpy.dtype[{int64}]]
+reveal_type(ar_iter[0, 0, 0])  # E: numpy.lib.arrayterator.Arrayterator[Any, numpy.dtype[{int64}]]
+reveal_type(ar_iter[..., 0, :])  # E: numpy.lib.arrayterator.Arrayterator[Any, numpy.dtype[{int64}]]