Подсказки типа Python и контекстные менеджеры

Как контекстный менеджер должен быть аннотирован подсказками типов Python?

import typing

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

В документации по contextlib documentation on contextlib много не упоминается о типах.

документация по typing.ContextManager не так уж полезна.

Также есть typing.Generator, в котором есть хотя бы пример. Означает ли это, что я должен использовать typing.Generator[None, None, None], а не typing.ContextManager?

import typing

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

Ответ 1

Всякий раз, когда я не уверен на 100%, какие типы принимает функция, я бы хотел обратиться к typeshed, которая является каноническим хранилищем подсказок типов для Python. Mypy, например, связывает и использует типизированные, чтобы помочь ему выполнить проверку типов, например.

Мы можем найти заглушки для contextlib здесь: https://github.com/python/typeshed/blob/master/stdlib/2and3/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 - мы можем снова проверить это сами , посоветовавшись с типами. Таким образом, если наша функция возвращает генератор, она все еще совместима с тем, что ожидает contextmanager, поэтому mypy принимает его без проблем.

Ответ 2

Версия Iterator[] не работает, если вы хотите вернуть ссылку на contextmanager. Например, следующий код:

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]:
    ...

Ответ 3

Возвращаемым типом функции, обернутой менеджером контекста, является Iterator[None].

from contextlib import contextmanager
from typing import Iterator

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

Ответ 4

С моим PyCharm я делаю следующее, чтобы заставить его подсказку работать:

from contextlib import contextmanager
from typing import ContextManager

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