可以对 numpy 数组/矩阵形状进行静态类型检查或 IDE 智能支持吗？

Question

是否可以为 numpy 数组/矩阵形状提供静态类型检查或 IDE 智能支持？

例如，如果我想象这样的事情：

A_MxN: NDArray(3,2) = ...
B_NxM: NDArray(2,3) = ...

更好的是：

N = 3
M = 2   
A_MxN: NDArray(M,N) = ...
B_NxM: NDArray(N,M) = ...

如果我将 A 分配给 B，我希望在开发期间（而不是运行时）得到 IDE 提示，表明形状不同。

类似：

A_MxN = B_NxM
Hint/Error: Declared shape 3,2 is not compatible with assigned shape 2,3

正如@simon提到的，这似乎是可能的：

M = Literal[3]
N = Literal[2]
A_MxN: np.ndarray[tuple[M,N], np.dtype[np.int32]]

但是如果我分配一个不满足形状要求的数组，linter 不会抛出错误。有人知道是否有像 mypy 或 Pyright 这样的类型检查器支持该功能？

Answer 1

现在似乎不可能，因为 nd.array 的构造函数会生成 NDArray[any]：https://github.com/microsoft/pyright/issues/8921