Type hint for numpy.ndarray containing unsignedinteger-CodePudding

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]

CodePudding user response：

You can use typing module from numpy:

import numpy as np
import numpy.typing as npt

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

CodePudding user response：

The NumPy documentation has you covered. You are looking for numpy.typing.NDArray:

from numpy import typing as npt

def myfunc(arr: npt.NDArray):
    pass