pythonnumpypython-typingmypypyright

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


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?


Solution

  • 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:

    It is thus necessary to manually cast the values:

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