8000 Merge pull request #27755 from jorenham/typing/ndarray-missing-annota… · numpy/numpy@44c2263 · GitHub
[go: up one dir, main page]

Skip to content

Commit 44c2263

Browse files
authored
Merge pull request #27755 from jorenham/typing/ndarray-missing-annotations
TYP: Annotate ``__setitem__``, ``__contains__`` and ``__iter__`` of ``ndarray``
2 parents e0eee7f + 2d02a3d commit 44c2263

File tree

1 file changed

+23
-20
lines changed

1 file changed

+23
-20
lines changed

numpy/__init__.pyi

Lines changed: 23 additions & 20 deletions
Original file line numberDiff line numberDiff line change
@@ -210,7 +210,7 @@ from typing import (
210210
# library include `typing_extensions` stubs:
211211
# https://github.com/python/typeshed/blob/main/stdlib/typing_extensions.pyi
212212
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
214214

215215
from numpy import (
216216
core,
@@ -1795,6 +1795,8 @@ _ArrayComplex_co: TypeAlias = NDArray[np.bool | integer[Any] | floating[Any] | c
17951795
_ArrayNumber_co: TypeAlias = NDArray[np.bool | number[Any]]
17961796
_ArrayTD64_co: TypeAlias = NDArray[np.bool | integer[Any] | timedelta64]
17971797

1798+
_ArrayIndexLike: TypeAlias = SupportsIndex | slice | EllipsisType | _ArrayLikeInt_co | None
1799+
17981800
# Introduce an alias for `dtype` to avoid naming conflicts.
17991801
_dtype: TypeAlias = dtype[_ScalarType]
18001802

@@ -1913,26 +1915,20 @@ class ndarray(_ArrayOrScalarCommon, Generic[_ShapeType_co, _DType_co]):
19131915
) -> ndarray[_ShapeType, _DType]: ...
19141916

19151917
@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: ...
19211921
@overload
1922-
def __getitem__(self, key: SupportsIndex | tuple[SupportsIndex, ...]) -> Any: ...
1922+
def __getitem__(self, key: _ArrayIndexLike | tuple[_ArrayIndexLike, ...], /) -> ndarray[_Shape, _DType_co]: ...
19231923
@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+
19321928
@overload
1933-
def __getitem__(self: NDArray[void], key: str) -> NDArray[Any]: ...
1929+
def __setitem__(self: NDArray[void], key: str | list[str], value: ArrayLike, /) -> None: ...
19341930
@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: ...
19361932

19371933
@property
19381934
def ctypes(self) -> _ctypes[int]: ...
@@ -2279,9 +2275,16 @@ class ndarray(_ArrayOrScalarCommon, Generic[_ShapeType_co, _DType_co]):
22792275
def __complex__(self: NDArray[number[Any] | np.bool | object_], /) -> complex: ...
22802276

22812277
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]: ...
22852288

22862289
# The last overload is for catching recursive objects whose
22872290
# nesting is too deep.

0 commit comments

Comments
 (0)
0