TST: Add typing tests for `nditer` and `nested_iters`

4 files changed, 62 insertions, 13 deletions

diff --git a/numpy/typing/tests/data/fail/multiarray.py b/numpy/typing/tests/data/fail/multiarray.py
index 603a0452b..22bcf8c92 100644
--- a/numpy/typing/tests/data/fail/multiarray.py
+++ b/numpy/typing/tests/data/fail/multiarray.py
@@ -47,3 +47,10 @@ np.datetime_as_string("2012")  # E: No overload variant
 np.compare_chararrays("a", b"a", "==", False)  # E: No overload variant
 
 np.add_docstring(func, None)  # E: incompatible type
+
+np.nested_iters([AR_i8, AR_i8])  # E: Missing positional argument
+np.nested_iters([AR_i8, AR_i8], 0)  # E: incompatible type
+np.nested_iters([AR_i8, AR_i8], [0])  # E: incompatible type
+np.nested_iters([AR_i8, AR_i8], [[0], [1]], flags=["test"])  # E: incompatible type
+np.nested_iters([AR_i8, AR_i8], [[0], [1]], op_flags=[["test"]])  # E: incompatible type
+np.nested_iters([AR_i8, AR_i8], [[0], [1]], buffersize=1.0)  # E: incompatible type
diff --git a/numpy/typing/tests/data/fail/nditer.py b/numpy/typing/tests/data/fail/nditer.py
new file mode 100644
index 000000000..1e8e37ee5
--- /dev/null
+++ b/numpy/typing/tests/data/fail/nditer.py
@@ -0,0 +1,8 @@
+import numpy as np
+
+class Test(np.nditer): ...  # E: Cannot inherit from final class
+
+np.nditer([0, 1], flags=["test"])  # E: incompatible type
+np.nditer([0, 1], op_flags=[["test"]])  # E: incompatible type
+np.nditer([0, 1], itershape=(1.0,))  # E: incompatible type
+np.nditer([0, 1], buffersize=1.0)  # E: incompatible type
diff --git a/numpy/typing/tests/data/reveal/multiarray.py b/numpy/typing/tests/data/reveal/multiarray.py
index cee51975e..ee818c08a 100644
--- a/numpy/typing/tests/data/reveal/multiarray.py
+++ b/numpy/typing/tests/data/reveal/multiarray.py
@@ -25,6 +25,8 @@ M: np.datetime64
 b_f8 = np.broadcast(AR_f8)
 b_i8_f8_f8 = np.broadcast(AR_i8, AR_f8, AR_f8)
 
+nditer_obj: np.nditer
+
 def func(a: int) -> bool: ...
 
 reveal_type(next(b_f8))  # E: tuple[Any]
@@ -123,3 +125,8 @@ reveal_type(np.compare_chararrays("a", "b", "!=", rstrip=False))  # E: numpy.nda
 reveal_type(np.compare_chararrays(b"a", b"a", "==", True))  # E: numpy.ndarray[Any, numpy.dtype[numpy.bool_]]
 
 reveal_type(np.add_docstring(func, "test"))  # E: None
+
+reveal_type(np.nested_iters([AR_i8, AR_i8], [[0], [1]], flags=["c_index"]))  # E: tuple[numpy.nditer]
+reveal_type(np.nested_iters([AR_i8, AR_i8], [[0], [1]], op_flags=[["readonly", "readonly"]]))  # E: tuple[numpy.nditer]
+reveal_type(np.nested_iters([AR_i8, AR_i8], [[0], [1]], op_dtypes=np.int_))  # E: tuple[numpy.nditer]
+reveal_type(np.nested_iters([AR_i8, AR_i8], [[0], [1]], order="C", casting="no"))  # E: tuple[numpy.nditer]
diff --git a/numpy/typing/tests/data/reveal/nditer.py b/numpy/typing/tests/data/reveal/nditer.py
index de2a4b5ed..473e922a2 100644
--- a/numpy/typing/tests/data/reveal/nditer.py
+++ b/numpy/typing/tests/data/reveal/nditer.py
@@ -1,19 +1,46 @@
-import copy
 import numpy as np
 
 nditer_obj: np.nditer
 
-with nditer_obj as context:
-    reveal_type(context)   # E: numpy.nditer
+reveal_type(np.nditer([0, 1], flags=["c_index"]))  # E: numpy.nditer
+reveal_type(np.nditer([0, 1], op_flags=[["readonly", "readonly"]]))  # E: numpy.nditer
+reveal_type(np.nditer([0, 1], op_dtypes=np.int_))  # E: numpy.nditer
+reveal_type(np.nditer([0, 1], order="C", casting="no"))  # E: numpy.nditer
 
-reveal_type(len(nditer_obj))  # E: builtins.int
-reveal_type(copy.copy(nditer_obj))  # E: numpy.nditer
-reveal_type(next(nditer_obj))  # E: Any
-reveal_type(iter(nditer_obj))  # E: typing.Iterator[Any]
-reveal_type(nditer_obj[1])  # E: Any
-reveal_type(nditer_obj[1:5])  # E: Any
+reveal_type(nditer_obj.dtypes)  # E: tuple[numpy.dtype[Any]]
+reveal_type(nditer_obj.finished)  # E: bool
+reveal_type(nditer_obj.has_delayed_bufalloc)  # E: bool
+reveal_type(nditer_obj.has_index)  # E: bool
+reveal_type(nditer_obj.has_multi_index)  # E: bool
+reveal_type(nditer_obj.index)  # E: int
+reveal_type(nditer_obj.iterationneedsapi)  # E: bool
+reveal_type(nditer_obj.iterindex)  # E: int
+reveal_type(nditer_obj.iterrange)  # E: tuple[builtins.int]
+reveal_type(nditer_obj.itersize)  # E: int
+reveal_type(nditer_obj.itviews)  # E: tuple[numpy.ndarray[Any, numpy.dtype[Any]]]
+reveal_type(nditer_obj.multi_index)  # E: tuple[builtins.int]
+reveal_type(nditer_obj.ndim)  # E: int
+reveal_type(nditer_obj.nop)  # E: int
+reveal_type(nditer_obj.operands)  # E: tuple[numpy.ndarray[Any, numpy.dtype[Any]]]
+reveal_type(nditer_obj.shape)  # E: tuple[builtins.int]
+reveal_type(nditer_obj.value)  # E: tuple[numpy.ndarray[Any, numpy.dtype[Any]]]
 
-nditer_obj[1] = 1
-nditer_obj[1:5] = 1
-del nditer_obj[1]
-del nditer_obj[1:5]
+reveal_type(nditer_obj.close())  # E: None
+reveal_type(nditer_obj.copy())  # E: numpy.nditer
+reveal_type(nditer_obj.debug_print())  # E: None
+reveal_type(nditer_obj.enable_external_loop())  # E: None
+reveal_type(nditer_obj.iternext())  # E: bool
+reveal_type(nditer_obj.remove_axis(0))  # E: None
+reveal_type(nditer_obj.remove_multi_index())  # E: None
+reveal_type(nditer_obj.reset())  # E: None
+
+reveal_type(len(nditer_obj))  # E: int
+reveal_type(iter(nditer_obj))  # E: Iterator[builtins.tuple[numpy.ndarray[Any, numpy.dtype[Any]]]]
+reveal_type(next(nditer_obj))  # E: tuple[numpy.ndarray[Any, numpy.dtype[Any]]]
+reveal_type(nditer_obj.__copy__())  # E: numpy.nditer
+with nditer_obj as f:
+    reveal_type(f)  # E: numpy.nditer
+reveal_type(nditer_obj[0])  # E: numpy.ndarray[Any, numpy.dtype[Any]]
+reveal_type(nditer_obj[:])  # E: tuple[numpy.ndarray[Any, numpy.dtype[Any]]]
+nditer_obj[0] = 0
+nditer_obj[:] = [0, 1]