Search code examples
pythonnumpymypy

Type hint for numpy.ndarray containing unsignedinteger

I have a numpy array that contains unsignedinteger, something like this:

arr = np.uint16([5, 100, 2000])

array([   5,  100, 2000], dtype=uint16)

This arr will be input to a function. I am wondering how the type hint of the function argument should look like?

def myfunc(arr: ?):
    pass

I was first thinking it should be arr: np.ndarray. But then mypy is complaining.
Argument "arr" to "myfunc" has incompatible type "unsignedinteger[_16Bit]"; expected "ndarray[Any, Any]" [arg-type]

Neither does arr: np.ndarray[np.uint16] work.
error: "ndarray" expects 2 type arguments, but 1 given [type-arg]

Solution

  • You can use typing module from numpy:

    import numpy as np
import numpy.typing as npt

def myfunc(arr: npt.NDArray[np.int16]):
    pass