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

那么，为什么按照 cmets 的建议使用 Generator[None, None, None] 也有效？

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

Answer 2

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

from contextlib import contextmanager
from typing import ContextManager

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

UPD：见下面的 cmets。看起来这件事让 PyC​​harm 开心，但 mypy 不开心

Answer 3

当您想要返回上下文管理器的引用时，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 4

基于PEP-585，正确的注释类型似乎是AbstractContextManager（参见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 5

上下文管理器包装的函数的返回类型是Iterator[None]。

from contextlib import contextmanager
from typing import Iterator

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

Answer 6

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

被修饰的函数在调用时必须返回一个generator-iterator

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)

Answer 7

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

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]注释实现解决了这个问题。