Specific type annotation for NumPy ndarray using mypy

Question

Support for type annotations was added in NumPy 1.20. I'm trying to figure out how to tell mypy that an array is filled with elements of a particular type, the annotation np.ndarray[np.dcomplex] gives the mypy error "ndarray" expects no type arguments, but 1 given.

EDIT: This question is different from Type hinting / annotation (PEP 484) for numpy.ndarray as that question was asked 4 years ago when there wasn't any official support for type hinting. I'm asking for what is the official way to do this, now that type hinting is actually natively supported by numpy 1.20. The documentation at https://numpy.org/doc/stable/reference/typing.html#module-numpy.typing that the top answer there points towards only seems to say things you shouldn't do with type hinting instead of explaining what you should be doing.

Answer 1

What you are looking for is the numpy.typing.NDArray class: https://numpy.org/doc/stable/reference/typing.html#numpy.typing.NDArray

numpy.typing.NDArray[A] is an alias for numpy.ndarray[Any, numpy.dtype[A]]:

import numpy as np
import numpy.typing as npt

a: npt.NDArray[np.complex64] = np.zeros((3, 3), dtype=np.complex64)
# reveal_type(a)  # -> numpy.ndarray[Any, numpy.dtype[numpy.complexfloating[numpy.typing._32Bit, numpy.typing._32Bit]]]
print(a)

prints

[[0.+0.j 0.+0.j 0.+0.j]
 [0.+0.j 0.+0.j 0.+0.j]
 [0.+0.j 0.+0.j 0.+0.j]]

Note that even though you annotate a as npt.NDArray[np.complex64], you still need to make sure that you pass the matching dtype to the factory on the right side.

a: npt.NDArray[np.complex64] = np.zeros((3, 3), dtype=np.float32)

passes the mypy check just as well.