Python 类型提示和上下文管理器

Question

上下文管理器应如何使用 Python 类型提示进行注释？

import typing

@contextlib.contextmanager
def foo() -> ???:
    yield

documentation on contextlib 没有太多提及类型。

documentation on typing.ContextManager 也不是很有帮助。

还有typing.Generator，至少有个例子。这是否意味着我应该使用 typing.Generator[None, None, None] 而不是 typing.ContextManager？

import typing

@contextlib.contextmanager
def foo() -> typing.Generator[None, None, None]:
    yield

Answer 1

每当我不能 100% 确定函数接受什么类型时，我喜欢查阅 typeshed，这是 Python 类型提示的规范存储库。例如，Mypy 直接捆绑并使用 typeshed 来帮助它执行类型检查。

我们可以在这里找到 contextlib 的存根：https://github.com/python/typeshed/blob/master/stdlib/contextlib.pyi

if sys.version_info >= (3, 2):
    class GeneratorContextManager(ContextManager[_T], Generic[_T]):
        def __call__(self, func: Callable[..., _T]) -> Callable[..., _T]: ...
    def contextmanager(func: Callable[..., Iterator[_T]]) -> Callable[..., GeneratorContextManager[_T]]: ...
else:
    def contextmanager(func: Callable[..., Iterator[_T]]) -> Callable[..., ContextManager[_T]]: ...

有点难听，不过我们关心的是这一行：

def contextmanager(func: Callable[..., Iterator[_T]]) -> Callable[..., ContextManager[_T]]: ...

它声明装饰器接受一个 Callable[..., Iterator[_T]] -- 一个带有任意参数并返回一些迭代器的函数。所以总而言之，这样做会很好：

@contextlib.contextmanager
def foo() -> Iterator[None]:
    yield

那么，为什么使用 Generator[None, None, None] 也像评论所建议的那样有效？

这是因为 Generator 是 Iterator 的子类型——我们可以自己再次检查一下 by consulting typeshed。所以，如果我们的函数 returns 是一个生成器，它仍然与 contextmanager 期望的兼容，所以 mypy 可以毫无问题地接受它。

Answer 2

当您想 return 上下文管理器的引用时，Iterator[] 版本不起作用。比如下面的代码：

from typing import Iterator

def assert_faster_than(seconds: float) -> Iterator[None]:
    return assert_timing(high=seconds)

@contextmanager
def assert_timing(low: float = 0, high: float = None) -> Iterator[None]:
    ...

将在 return assert_timing(high=seconds) 行产生错误：

Incompatible return value type (got "_GeneratorContextManager[None]", expected "Iterator[None]")

函数的任何合法用法：

with assert_faster_than(1):
    be_quick()

会产生这样的结果：

"Iterator[None]" has no attribute "__enter__"; maybe "__iter__"?
"Iterator[None]" has no attribute "__exit__"; maybe "__next__"?
"Iterator[None]" has no attribute "__enter__"; maybe "__iter__"?
"Iterator[None]" has no attribute "__exit__"; maybe "__next__"?

你可以这样修复它...

def assert_faster_than(...) -> Iterator[None]:
    with assert_timing(...):
        yield

但我将使用新的 ContextManager[] 对象来代替装饰器的 mypy:

from typing import ContextManager

def assert_faster_than(seconds: float) -> ContextManager[None]:
    return assert_timing(high=seconds)

@contextmanager  # type: ignore
def assert_timing(low: float = 0, high: float = None) -> ContextManager[None]:
    ...

Answer 3

一个。 @contextmanager修饰的函数的return类型是Iterator[None].

from contextlib import contextmanager
from typing import Iterator

@contextmanager
def foo() -> Iterator[None]:
    yield

乙。上下文管理器本身的类型是 AbstractContextManager:

from contextlib import AbstractContextManager

def make_it_so(context: AbstractContextManager) -> None:
    with context:
        ...

您可能还会看到使用了 typing.ContextManager，但自 Python 3.9 以来 deprecated 支持 contextlib.AbstractContextManager。

Answer 4

对于我的 PyCharm，我执行以下操作以使其类型提示起作用：

from contextlib import contextmanager
from typing import ContextManager

@contextmanager
def session() -> ContextManager[Session]:
    yield Session(...)

UPD：请参阅下面的评论。看起来这个东西让 PyCharm 快乐，但不是 mypy

Answer 5

基于PEP-585 the correct annotation type seems to be AbstractContextManager (see https://www.python.org/dev/peps/pep-0585/#implementation）。您可以使用以下代码：

import contextlib

@contextlib.contextmanager
def foo() -> contextlib.AbstractContextManager[None]:
    yield

这是唯一可以与 PyCharm（以及 typing.ContextManager 一起正确工作的解决方案，但应该从 Python 3.9 开始弃用该解决方案。当您在 with 语句（类型提示）中使用它时，它会正确地帮助您，这非常有帮助。

但是当我回到最初的问题时（“上下文管理器应该如何用 Python 类型提示进行注释？”）这取决于。从我的角度来看，正确的应该是我提到的那个。但这似乎不适用于 mypy（还）。有一些关于此 PEP 的更新（请参阅 https://github.com/python/mypy/issues/7907），但由于我对 mypy 的经验不多，因此我可能会遗漏一些内容。

Answer 6

我在实现抽象方法的时候遇到了类似的问题：

class Abstract(ABC):
    @abstractmethod
    def manager(self) -> ContextManager[None]:
        pass


class Concrete(Abstract):
    @contextmanager
    def manager(self) -> Iterator[None]:
        try:
            yield
        finally:
            pass

用 ContextManager[None] 注释抽象方法并用 Iterator[None] 注释实现解决了问题。

Answer 7

我在这里找不到关于注释上下文管理器的好答案，这些上下文管理器以通过 Python 3.10 下的 mypy 检查的方式产生值。根据 Python 3.10 documentation for contextlib.contextmanager

The function being decorated must return a generator-iterator when called

typing.Generators被注释为Generator[YieldType, SendType, ReturnType]。因此，对于产生 pathlib.Path 的函数，我们可以这样注释我们的函数：

from typing import Generator
from contextlib import contextmanager

@contextmanager
def working_directory() -> Generator[Path, None, None]:
    with TemporaryDirectory() as td:
        yield Path(td)

但是，未指定 SendType 或 ReturnType 的 Generators 可以改为注释为 typing.Iterator:

from typing import Iterator
from contextlib import contextmanager

@contextmanager
def working_directory() -> Iterator[Path]:
    with TemporaryDirectory() as td:
        yield Path(td)

最后，由于 PEP 585 -- Type Hinting Generics In Standard Collections 在 Python 3.9 中被采用，typing.Iterator 和 typing.Generator 被弃用，取而代之的是 collections.abc 实现

from collections.abc import Iterator
from contextlib import contextmanager

@contextmanager
def working_directory() -> Iterator[Path]:
    with TemporaryDirectory() as td:
        yield Path(td)

Python 类型提示和上下文管理器

Python type hints and context managers

python

mypy