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: ?):

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]

Asked By: Andi



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

from numpy import typing as npt

def myfunc(arr: npt.NDArray):

Answered By: JustLearning

You can use typing module from numpy:

import numpy as np
import numpy.typing as npt

def myfunc(arr: npt.NDArray[np.int16]):
Answered By: Corralien
Categories: questions Tags: , ,
Answers are sorted by their score. The answer accepted by the question owner as the best is marked with
at the top-right corner.