Can Pyright/MyPy deduce the type of an entry of an ndarray?

Question

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?

Answer 1

This is not currently possible.

As defined in the type stubs, ndarray.__getitem__ has 5 @overloads. 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:

• Pyright only reads information out of type stubs, so it's not possible to modify its behaviour.
• Mypy allows for plugins, and Numpy bundles its own Mypy plugin, but this plugin doesn't seem to have the inference capabilities you are looking for.

It is thus necessary to manually cast the values:

from typing import cast

y = cast(int, x[0])
reveal_type(y)  # int