来自 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 。

有关更多详细信息，请参阅此问题和此问题。

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

一个例子：

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 可以很好地检测到并正确运行：