上下文管理器的类型提示

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

您的函数实际上并不是 returning 一个元组。相反，它是 yielding 一个——如果你在没有上下文管理器的情况下调用 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)

您可以找到有关生成器 in the mypy docs 的更多信息，但简而言之，它接受三种类型：生成器产生的任何类型、生成器可以发送的任何值的类型以及生成器的类型不管你的发电机最终值是多少 returns。我们这里不关心后两者，所以我们将它们保留为 None.

虽然这种类型最终会让人感觉笨拙。更简洁的替代方法是将 return 类型注释为迭代器：

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)

之所以可行，是因为所有生成器实际上都是迭代器 -- 生成器是迭代器的子类型。所以选择更通用的return类型就好了。

上下文管理器的类型提示

Type hints for context manager

type-hinting

contextmanager

python-3.x