python - 来自 plt.subplots() 的 matplotlib 轴的精确类型注释数组 (numpy.ndarray)

Question

我希望在使用 VSCode Pylance 类型检查器时没有错误。

如何在以下代码中正确输入axs:

import matplotlib.pyplot as plt
fig, axs = plt.subplots(2, 2)

在下图中，您可以看到 VSCode 上的 Pylance 正在检测到错误。

Answer 1

事实证明，强类型化 axs 变量根本不简单，需要很好地理解如何类型化 np.ndarray。

参见this question和 this question了解更多详情。

最简单、最强大的解决方案是用 ' 字符包装 numpy.ndarray，以避免臭名昭著的 TypeError: 'numpy._DTypeMeta' object is not subscriptable当Python尝试解释表达式中的[]时。

一个例子:

import matplotlib.pyplot as plt
import numpy as np
import numpy.typing as npt
import seaborn as sns
from typing import cast, Type, Sequence
import typing 

sns.set() 

# Some example data to display
x = np.linspace(0, 2 * np.pi, 400)
y = np.sin(x ** 2)

fig, axs = plt.subplots(
    2, 2, 
    figsize=(12, 10) # set graph size
)

# typechecking operation
NDArrayOfAxes: typing.TypeAlias = 'np.ndarray[Sequence[Sequence[plt.Axes]], np.dtype[np.object_]]'
axs = cast(np.ndarray, axs)

axs[0, 0].plot(x, y)
axs[0, 0].set_title("main")
axs[1, 0].plot(x, y**2)
axs[1, 0].set_title("shares x with main")
axs[1, 0].sharex(axs[0, 0])
axs[0, 1].plot(x + 1, y + 1)
axs[0, 1].set_title("unrelated")
axs[1, 1].plot(x + 2, y + 2)
axs[1, 1].set_title("also unrelated")
fig.tight_layout()

Pylance 可以很好地检测到它并正确运行: