TYP: Add method annotations in ndarray

Author: jorenham

numpy/__init__.pyi

@@ -210,7 +210,7 @@ from typing import (
 # library include `typing_extensions` stubs:
 # https://github.com/python/typeshed/blob/main/stdlib/typing_extensions.pyi
 from _typeshed import StrOrBytesPath, SupportsFlush, SupportsLenAndGetItem, SupportsWrite
-from typing_extensions import CapsuleType, Generic, LiteralString, Protocol, Self, TypeVar, deprecated, overload
+from typing_extensions import CapsuleType, Generic, LiteralString, Protocol, Self, TypeVar, Unpack, deprecated, overload
 
 from numpy import (
     core,
@@ -1792,6 +1792,8 @@ _ArrayComplex_co: TypeAlias = NDArray[np.bool | integer[Any] | floating[Any] | c
 _ArrayNumber_co: TypeAlias = NDArray[np.bool | number[Any]]
 _ArrayTD64_co: TypeAlias = NDArray[np.bool | integer[Any] | timedelta64]
 
+_ArrayIndexLike: TypeAlias = SupportsIndex | slice | EllipsisType | _ArrayLikeInt_co | None
+
 # Introduce an alias for `dtype` to avoid naming conflicts.
 _dtype: TypeAlias = dtype[_ScalarType]
 
@@ -1906,26 +1908,20 @@ class ndarray(_ArrayOrScalarCommon, Generic[_ShapeType_co, _DType_co]):
     ) -> ndarray[_ShapeType, _DType]: ...
 
     @overload
-    def __getitem__(self, key: (
-        NDArray[integer[Any]]
-        | NDArray[np.bool]
-        | tuple[NDArray[integer[Any]] | NDArray[np.bool], ...]
-    )) -> ndarray[_Shape, _DType_co]: ...
+    def __getitem__(self, key: _ArrayInt_co | tuple[_ArrayInt_co, ...], /) -> ndarray[_Shape, _DType_co]: ...
+    @overload
+    def __getitem__(self, key: SupportsIndex | tuple[SupportsIndex, ...], /) -> Any: ...
     @overload
-    def __getitem__(self, key: SupportsIndex | tuple[SupportsIndex, ...]) -> Any: ...
+    def __getitem__(self, key: _ArrayIndexLike | tuple[_ArrayIndexLike, ...], /) -> ndarray[_Shape, _DType_co]: ...
     @overload
-    def __getitem__(self, key: (
-        None
-        | slice
-        | EllipsisType
-        | SupportsIndex
-        | _ArrayLikeInt_co
-        | tuple[None | slice | EllipsisType | _ArrayLikeInt_co | SupportsIndex, ...]
-    )) -> ndarray[_Shape, _DType_co]: ...
+    def __getitem__(self: NDArray[void], key: str, /) -> ndarray[_ShapeType_co, np.dtype[Any]]: ...
+    @overload
+    def __getitem__(self: NDArray[void], key: list[str], /) -> ndarray[_ShapeType_co, _dtype[void]]: ...
+
     @overload
-    def __getitem__(self: NDArray[void], key: str) -> NDArray[Any]: ...
+    def __setitem__(self: NDArray[void], key: str | list[str], value: ArrayLike, /) -> None: ...
     @overload
-    def __getitem__(self: NDArray[void], key: list[str]) -> ndarray[_ShapeType_co, _dtype[void]]: ...
+    def __setitem__(self, key: _ArrayIndexLike | tuple[_ArrayIndexLike, ...], value: ArrayLike, /) -> None: ...
 
     @property
     def ctypes(self) -> _ctypes[int]: ...
@@ -2272,9 +2268,16 @@ class ndarray(_ArrayOrScalarCommon, Generic[_ShapeType_co, _DType_co]):
     def __complex__(self: NDArray[number[Any] | np.bool | object_], /) -> complex: ...
 
     def __len__(self) -> int: ...
-    def __setitem__(self, key, value): ...
-    def __iter__(self) -> Any: ...
-    def __contains__(self, key) -> builtins.bool: ...
+    def __contains__(self, value: object, /) -> builtins.bool: ...
+
+    @overload  # == 1-d & object_
+    def __iter__(self: ndarray[tuple[int], dtype[object_]], /) -> Iterator[Any]: ...
+    @overload  # == 1-d
+    def __iter__(self: ndarray[tuple[int], dtype[_SCT]], /) -> Iterator[_SCT]: ...
+    @overload  # >= 2-d
+    def __iter__(self: ndarray[tuple[int, int, Unpack[tuple[int, ...]]], dtype[_SCT]], /) -> Iterator[NDArray[_SCT]]: ...
+    @overload  # ?-d
+    def __iter__(self, /) -> Iterator[Any]: ...
 
     # The last overload is for catching recursive objects whose
     # nesting is too deep.