plt.subplots()中matplotlib轴的精确类型注解数组(numpy.ndarray)

5us2dqdw  于 2022-11-15  发布在  其他
关注(0)|答案(2)|浏览(149)

我希望在使用VSCode Pylance类型检查器时没有错误。
如何在下面的代码中正确键入axs

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

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

vwhgwdsa

vwhgwdsa1#

事实证明,强式输入axs变量并不简单,需要很好地理解如何输入np.ndarray
请参阅本题和本题了解更多详细信息。
最简单、最强大的解决方案是用'字符 Package numpy.ndarray,以避免臭名昭著的TypeError:当Python试图解释表达式中的[]时,'numpy._DTypeMeta'对象是不可订阅的。
举个例子:

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能够很好地检测并正确运行:

c9x0cxw0

c9x0cxw02#

由于我不打算理解的原因,currently accepted answer在我的设置中不断给出警告、错误,并且缺乏成功的类型推理。
我所做的是这样的,它允许切片,并允许Pylance理解.plotAxes.plot

from typing import TypeVar, Generic

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

# Little helper class, which is only used as a type.
DType = TypeVar("DType")
class Array(np.ndarray, Generic[DType]):
    def __getitem__(self, key) -> DType:
        return super().__getitem__(key)

# Force assign the type, which is correct for most intents and purposes
fig, axs_ = plt.subplots(2, 2)
axs: Array[Axes] = axs_ # type: ignore

# Use as an ndarray of Axes
axs[0,0].plot(...)

相关问题