Merge pull request numpy#19554 from MatthieuDartiailh/nditer-typing

MAINT: add missing dunder method to nditer type hints

Author: mattip

numpy/__init__.pyi

@@ -173,6 +173,7 @@ from typing import (
     Generic,
     IO,
     Iterable,
+    Iterator,
     List,
     Mapping,
     NoReturn,
@@ -780,6 +781,21 @@ class nditer:
         buffersize: Any = ...,
     ) -> Any: ...
     def __getattr__(self, key: str) -> Any: ...
+    def __enter__(self) -> nditer: ...
+    def __exit__(
+        self,
+        exc_type: None | Type[BaseException],
+        exc_value: None | BaseException,
+        traceback: None | TracebackType,
+    ) -> None: ...
+    def __iter__(self) -> Iterator[Any]: ...
+    def __next__(self) -> Any: ...
+    def __len__(self) -> int: ...
+    def __copy__(self) -> nditer: ...
+    def __getitem__(self, index: SupportsIndex | slice) -> Any: ...
+    def __setitem__(self, index: SupportsIndex | slice, value: Any) -> None: ...
+    def __delitem__(self, key: SupportsIndex | slice) -> None: ...
+
 
 class poly1d:
     def __init__(

numpy/typing/tests/data/reveal/nditer.py

@@ -0,0 +1,19 @@
+import copy
+import numpy as np
+
+nditer_obj: np.nditer
+
+with nditer_obj as context:
+    reveal_type(context)   # 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
+
+nditer_obj[1] = 1
+nditer_obj[1:5] = 1
+del nditer_obj[1]
+del nditer_obj[1:5]