來自 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 可以很好地檢測到並正確運行：