Specific type annotation for NumPy ndarray using mypy
Answer a 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.
Answers
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.
更多推荐

所有评论(0)