python - Type hint for NumPy ndarray dtype?

Question

I would like a function to include a type hint for NumPy ndarray's alongside with its dtype.

With lists, for example, one could do the following...

def foo(bar: List[int]):
   ...

...in order to give a type hint that bar has to be list consisting of int's.

Unfortunately, this syntax throws exceptions for NumPy ndarray:

def foo(bar: np.ndarray[np.bool]):
   ...

> np.ndarray[np.bool]) (...) TypeError: 'type' object is not subscriptable

Is it possible to give dtype-specific type hints for np.ndarray?

Answer 1

You could check out nptyping:

from nptyping import NDArray, Bool

def foo(bar: NDArray[Bool]):
   ...

Or you could just use strings for type hints:

def foo(bar: 'np.ndarray[np.bool]'):
   ...