上下文管理器的类型提示

Question

我想要一个 pyplot 图形的上下文管理器，基本上像这样：

from contextlib import contextmanager
import matplotlib.pyplot as plt

@contextmanager
def subplots():
  (fig, ax) = plt.subplots()
  try:
    yield (fig, ax)
  finally:
    plt.close(fig)

是否可以对返回的元组实现类型提示？天真的

import typing
def subplots() -> typing.Tuple[plt.Figure, plt.Axes]

不起作用。

Answer 1

您的函数实际上并没有返回元组。相反，它会产生一个 - 如果您subplots()在没有上下文管理器的情况下调用，您将返回一个生成器对象。因此，您需要对其进行注释：

from typing import Tuple, Generator
from contextlib import contextmanager
import matplotlib.pyplot as plt

@contextmanager
def subplots() -> Generator[Tuple[plt.Figure, plt.Axes], None, None]:
  (fig, ax) = plt.subplots()
  try:
    yield (fig, ax)
  finally:
    plt.close(fig)

您可以在 mypy 文档 中找到有关 Generator 的更多信息，但简而言之，它接受三种类型：生成器生成的任何值的类型、生成器可以发送的任何值的类型以及生成器最终返回的任何值的类型。我们在这里不关心后两者，因此我们将它们保留为 None。

不过，这种类型最终确实会让人感觉笨拙。更简洁的替代方法是将返回类型注释为迭代器：

from typing import Tuple, Iterator
from contextlib import contextmanager
import matplotlib.pyplot as plt

@contextmanager
def subplots() -> Iterator[Tuple[plt.Figure, plt.Axes]]:
  (fig, ax) = plt.subplots()
  try:
    yield (fig, ax)
  finally:
    plt.close(fig)

这是可行的，因为所有生成器实际上都是迭代器——生成器是迭代器的子类型。所以选择更通用的返回类型就可以了。