Specific type annotation for NumPy ndarray using mypy

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.

Leave a Comment

Hata!: SQLSTATE[HY000] [1045] Access denied for user 'divattrend_liink'@'localhost' (using password: YES)