How to narrow numpy array dtype in if statements

[gaufde]

Hello!

I am using some scientific python packages that don't have great type hints. I want to check that I am receiving the type of numpy array I expect before feeding the result to my own functions. However, I can't figure out how to get Pyright to understand this. For example:

from typing import Union
import numpy as np
import numpy.typing as npt


def func_1(arr: npt.NDArray[np.int_]) -> npt.NDArray[Union[np.float64, np.int_]]:
    return arr.astype(np.int64)


def func_2(arr: npt.NDArray[np.int_]) -> npt.NDArray[np.int_]:
    return arr


arr = np.ones((10,), dtype=np.int64)

arr = func_1(arr)

is_correct_type = arr.dtype == np.int_ or np.issubdtype(arr.dtype, np.int_)

if is_correct_type:
    arr = func_2(arr)

However, Pyright still throws an error:

Argument of type "NDArray[float64 | int_]" cannot be assigned to parameter "arr" of type "NDArray[int_]" in function "func_2"
"ndarray[Any, dtype[float64 | int_]]" is incompatible with "ndarray[Any, dtype[int_]]"
Type parameter "DType_co@ndarray" is covariant, but "dtype[float64 | int]" is not a subtype of "dtype[int_]"
"dtype[float64 | int_]" is incompatible with "dtype[int_]"
Type parameter "DTypeScalar_co@dtype" is covariant, but "float64 | int" is not a subtype of "int_"
Type "float64 | int_" is incompatible with type "int_" (reportArgumentType)

This error is corroborated by the hover information in my IDE. However, when I debug the code is_correct_type is True

[erictraut]

The Python static type system isn't rich enough to express the notion that "this object has an attribute called dtype that contains an immutable value that relates to the type of class". A static type checker therefore cannot narrow the type of arr in your example without additional information.

The TypeGuard type (introduced in PEP 647) was designed to handle cases like this.

from typing import Any, TypeGuard, TypeVar

T = TypeVar("T", bound=np.generic)

def is_correct_type(arr: npt.NDArray[Any], typ: type[T]) -> TypeGuard[npt.NDArray[T]]:
    return arr.dtype == np.int_ or np.issubdtype(arr.dtype, np.int_)

if is_correct_type(arr, np.int_):
    arr = func_2(arr)

Or if you're using Python 3.12, you can use the new generics syntax:

def is_correct_type[T: np.generic](arr: npt.NDArray[Any], typ: type[T]) -> TypeGuard[npt.NDArray[T]]:
    return arr.dtype == np.int_ or np.issubdtype(arr.dtype, np.int_)

[gaufde]

Thanks for the excellent explanation! This taught me so much about typing in Python.