python-3.x - 使用 mypy 对 NumPy ndarray 进行特定类型注释

Question

NumPy 1.20 中添加了对类型注释的支持。
我试图弄清楚如何告诉 mypy 一个数组填充了特定类型的元素，注释 np.ndarray[np.dcomplex]给出 mypy 错误 "ndarray" expects no type arguments, but 1 given .
编辑:这个问题不同于 Type hinting / annotation (PEP 484) for numpy.ndarray因为这个问题是在 4 年前提出的，当时没有任何官方支持类型提示。我问的是什么是官方的方法，现在numpy 1.20实际上支持类型提示。文档位于 https://numpy.org/doc/stable/reference/typing.html#module-numpy.typing那里的最高答案似乎只是说你不应该用类型提示做的事情，而不是解释你应该做什么。

Answer 1

您要找的是numpy.typing.NDArray类(class):https://numpy.org/doc/stable/reference/typing.html#numpy.typing.NDArraynumpy.typing.NDArray[A]是 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)

打印

[[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]]

请注意，即使您注释了 a如 npt.NDArray[np.complex64] ，你仍然需要确保你通过了匹配的 dtype到右边的工厂。

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

同样通过 mypy 检查。