@@ -210,7 +210,7 @@ from typing import (
210
210
# library include `typing_extensions` stubs:
211
211
# https://github.com/python/typeshed/blob/main/stdlib/typing_extensions.pyi
212
212
from _typeshed import StrOrBytesPath , SupportsFlush , SupportsLenAndGetItem , SupportsWrite
213
- from typing_extensions import CapsuleType , Generic , LiteralString , Protocol , Self , TypeVar , deprecated , overload
213
+ from typing_extensions import CapsuleType , Generic , LiteralString , Protocol , Self , TypeVar , Unpack , deprecated , overload
214
214
215
215
from numpy import (
216
216
core ,
@@ -1795,6 +1795,8 @@ _ArrayComplex_co: TypeAlias = NDArray[np.bool | integer[Any] | floating[Any] | c
1795
1795
_ArrayNumber_co : TypeAlias = NDArray [np .bool | number [Any ]]
1796
1796
_ArrayTD64_co : TypeAlias = NDArray [np .bool | integer [Any ] | timedelta64 ]
1797
1797
1798
+ _ArrayIndexLike : TypeAlias = SupportsIndex | slice | EllipsisType | _ArrayLikeInt_co | None
1799
+
1798
1800
# Introduce an alias for `dtype` to avoid naming conflicts.
1799
1801
_dtype : TypeAlias = dtype [_ScalarType ]
1800
1802
@@ -1913,26 +1915,20 @@ class ndarray(_ArrayOrScalarCommon, Generic[_ShapeType_co, _DType_co]):
1913
1915
) -> ndarray [_ShapeType , _DType ]: ...
1914
1916
1915
1917
@overload
1916
- def __getitem__ (self , key : (
1917
- NDArray [integer [Any ]]
1918
- | NDArray [np .bool ]
1919
- | tuple [NDArray [integer [Any ]] | NDArray [np .bool ], ...]
1920
- )) -> ndarray [_Shape , _DType_co ]: ...
1918
+ def __getitem__ (self , key : _ArrayInt_co | tuple [_ArrayInt_co , ...], / ) -> ndarray [_Shape , _DType_co ]: ...
1919
+ @overload
1920
+ def __getitem__ (self , key : SupportsIndex | tuple [SupportsIndex , ...], / ) -> Any : ...
1921
1921
@overload
1922
- def __getitem__ (self , key : SupportsIndex | tuple [SupportsIndex , ...]) -> Any : ...
1922
+ def __getitem__ (self , key : _ArrayIndexLike | tuple [_ArrayIndexLike , ...], / ) -> ndarray [ _Shape , _DType_co ] : ...
1923
1923
@overload
1924
- def __getitem__ (self , key : (
1925
- None
1926
- | slice
1927
- | EllipsisType
1928
- | SupportsIndex
1929
- | _ArrayLikeInt_co
1930
- | tuple [None | slice | EllipsisType | _ArrayLikeInt_co | SupportsIndex , ...]
1931
- )) -> ndarray [_Shape , _DType_co ]: ...
1924
+ def __getitem__ (self : NDArray [void ], key : str , / ) -> ndarray [_ShapeType_co , np .dtype [Any ]]: ...
1925
+ @overload
1926
+ def __getitem__ (self : NDArray [void ], key : list [str ], / ) -> ndarray [_ShapeType_co , _dtype [void ]]: ...
1927
+
1932
1928
@overload
1933
- def __getitem__ (self : NDArray [void ], key : str ) -> NDArray [ Any ] : ...
1929
+ def __setitem__ (self : NDArray [void ], key : str | list [ str ], value : ArrayLike , / ) -> None : ...
1934
1930
@overload
1935
- def __getitem__ (self : NDArray [ void ], key : list [ str ] ) -> ndarray [ _ShapeType_co , _dtype [ void ]] : ...
1931
+ def __setitem__ (self , key : _ArrayIndexLike | tuple [ _ArrayIndexLike , ... ], value : ArrayLike , / ) -> None : ...
1936
1932
1937
1933
@property
1938
1934
def ctypes (self ) -> _ctypes [int ]: ...
@@ -2279,9 +2275,16 @@ class ndarray(_ArrayOrScalarCommon, Generic[_ShapeType_co, _DType_co]):
2279
2275
def __complex__ (self : NDArray [number [Any ] | np .bool | object_ ], / ) -> complex : ...
2280
2276
2281
2277
def __len__ (self ) -> int : ...
2282
- def __setitem__ (self , key , value ): ...
2283
- def __iter__ (self ) -> Any : ...
2284
- def __contains__ (self , key ) -> builtins .bool : ...
2278
+ def __contains__ (self , value : object , / ) -> builtins .bool : ...
2279
+
2280
+ @overload # == 1-d & object_
2281
+ def __iter__ (self : ndarray [tuple [int ], dtype [object_ ]], / ) -> Iterator [Any ]: ...
2282
+ @overload # == 1-d
2283
+ def __iter__ (self : ndarray [tuple [int ], dtype [_SCT ]], / ) -> Iterator [_SCT ]: ...
2284
+ @overload # >= 2-d
2285
+ def __iter__ (self : ndarray [tuple [int , int , Unpack [tuple [int , ...]]], dtype [_SCT ]], / ) -> Iterator [NDArray [_SCT ]]: ...
2286
+ @overload # ?-d
2287
+ def __iter__ (self , / ) -> Iterator [Any ]: ...
2285
2288
2286
2289
# The last overload is for catching recursive objects whose
2287
2290
# nesting is too deep.
0 commit comments