How can I annotate an ndarray
so that Pyright/Mypy Intellisense can deduce the type of an entry? What can I fill in for ???
in
x: ??? = np.array([1, 2, 3], dtype=int)
so that
y = x[0]
is identified as an integer as rather than Any
?
This is not currently possible.
As defined in the type stubs, ndarray.__getitem__
has 5 @overload
s. When the type checker evaluates x[0]
, it uses the second one, which returns Any
:
@overload
def __getitem__(self, key: SupportsIndex | tuple[SupportsIndex, ...]) -> Any: ...
The other four all return ndarray
/NDArray
:
@overload
def __getitem__(self, key: (
NDArray[integer[Any]]
| NDArray[np.bool]
| tuple[NDArray[integer[Any]] | NDArray[np.bool], ...]
)) -> 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]: ...
@overload
def __getitem__(self: NDArray[void], key: str) -> NDArray[Any]: ...
@overload
def __getitem__(self: NDArray[void], key: list[str]) -> ndarray[_ShapeType_co, _dtype[void]]: ...
This means that, as long as x
is determined to be of type ndarray
, x[0]
cannot be inferred as int
.
As for type checkers' support:
It is thus necessary to manually cast the values:
from typing import cast
y = cast(int, x[0])
reveal_type(y) # int