使用 mypy 的 NumPy ndarray 的特定类型注释

Posted 2023-03-22

【中文标题】使用 mypy 的 NumPy ndarray 的特定类型注释

【英文标题】：Specific type annotation for NumPy ndarray using mypy

【发布时间】：2021-05-26 16:12:22

【问题描述】：

在 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 年前提出的，当时没有任何官方支持类型提示。我在问什么是 official 方法来做到这一点，现在类型提示实际上是由 numpy 1.20 原生支持的。 https://numpy.org/doc/stable/reference/typing.html#module-numpy.typing 上的文档，那里的最佳答案似乎只说明了你不应该用类型提示做的事情，而不是解释你应该做什么。

【问题讨论】：

这能回答你的问题吗？ Type hinting / annotation (PEP 484) for numpy.ndarray

那个答案指向没有解释我想知道的事情的文档。文档似乎只是说你不应该做的事情，而不是你应该做的事情。

请注意，对场外资源的请求是题外话。如果您要特别地询问如何注释为dtype数组，则无需询问“有人指向教程”；你可能想edit你的问题来删除这个请求。

我认为简短的回答是你有点太早了。当然 [来源][1] 引用了一个通用的 ndarray。但是，在最新的 numpy (1.20.1) 的全新 pip 安装中查看可比较的文件，它丢失了。实际上，您的选择是等待它正确发布，或者从 master 构建您自己的版本。无论哪种情况，您都需要注释np.ndarray[Any, np.dcomplex]，因为第一个 Type 参数实际上是用于形状的。 [1]：github.com/numpy/numpy/blob/main/numpy/__init__.pyi

【参考方案1】：

您正在寻找的是numpy.typing.NDArray 类：https://numpy.org/doc/stable/reference/typing.html#numpy.typing.NDArray

numpy.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 检查。