如何正确标注plt.subplots返回的axs类型消除Pylance报错

plt.subplots()的返回类型会根据传入的行列参数动态变化，当传入的nrows、ncols均大于1时，返回的axs是存储Axes实例的numpy多维数组，Pylance默认推断时经常无法识别具体类型导致报错，按下面的方法加标注即可解决。

• 首先导入标注需要的两个类型：
  • matplotlib.axes.Axes：对应数组内存储的坐标轴实例类型
  • numpy.typing.NDArray：numpy官方提供的数组泛型标注类型
• 给解构得到的axs变量追加类型注解，完整可过类型检查的代码如下：

import matplotlib.pyplot as plt
from matplotlib.axes import Axes
from numpy.typing import NDArray

fig, axs = plt.subplots(2, 2)
axs: NDArray[Axes]

加完标注后，Pylance可以正确识别axs[i, j]的返回值为Axes类型，后续调用.plot()、.set_title()等坐标轴方法时，会正常触发代码补全和参数校验，不会再报类型错误。

低版本兼容写法 #

如果你的numpy版本低于1.21，不存在numpy.typing.NDArray类型，可以换用numpy原生的数组标注写法，不需要额外导入NDArray，若提示Any未定义，从typing模块导入即可：

from typing import Any
import numpy as np
import matplotlib.pyplot as plt
from matplotlib.axes import Axes

fig, axs = plt.subplots(2, 2)
axs: np.ndarray[Any, np.dtype[Axes]]

特殊场景说明 #

如果调用plt.subplots()时传入的nrows、ncols均为1，返回的axs是单个Axes实例而非数组，直接标注为axs: Axes即可。

内容的提问来源于stack exchange，提问作者Onyr