Python type hints and context managers

Question

How should a context manager be annotated with Python type hints?

import typing

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

The documentation on contextlib doesn't mention types much.

The documentation on typing.ContextManager is not all that helpful either.

There's also typing.Generator, which at least has an example. Does that mean I should use typing.Generator[None, None, None] and not typing.ContextManager?

import typing

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

Answer 1

Whenever I'm not 100% sure what types a function accepts, I like to consult typeshed, which is the canonical repository of type hints for Python. Mypy directly bundles and uses typeshed to help it perform its typechecking, for example.

We can find the stubs for contextlib here: 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]]: ...

It's a little overwhelming, but the line we care about is this one:

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

It states that the decorator takes in a Callable[..., Iterator[_T]] -- a function with arbitrary arguments returning some iterator. So in conclusion, it would be fine to do:

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

So, why does using Generator[None, None, None] also work, as suggested by the comments?

It's because Generator is a subtype of Iterator -- we can again check this for ourselves by consulting typeshed. So, if our function returns a generator, it's still compatible with what contextmanager expects so mypy accepts it without an issue.

Answer 2

With my PyCharm, I do the following to make its type hinting work:

from contextlib import contextmanager
from typing import ContextManager

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

UPD: see comments below. Looks like this thing makes PyCharm happy, but not mypy

Answer 3

I didn't find a good answer here around annotating contextmanagers which yield values in a way which passes mypy checks under Python 3.10. According to the Python 3.10 documentation for contextlib.contextmanager

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

typing.Generators are annotated as Generator[YieldType, SendType, ReturnType]. So, in the case of a function which yields a pathlib.Path, we can annotate our functions like this:

from typing import Generator
from contextlib import contextmanager

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

However, Generators which don't specify SendType or ReturnType can instead be annotated as typing.Iterator:

from typing import Iterator
from contextlib import contextmanager

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

Finally, since PEP 585 -- Type Hinting Generics In Standard Collections was adopted in Python 3.9, typing.Iterator and typing.Generator are deprecated in favour of the collections.abc implementations

from collections.abc import Iterator
from contextlib import contextmanager

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

Answer 4

A. The return type of a function decorated by @contextmanager is Iterator[None].

from contextlib import contextmanager
from typing import Iterator

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

B. The type of the context manager itself is AbstractContextManager:

from contextlib import AbstractContextManager

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

You may also see typing.ContextManager used, but that has been deprecated in favor of contextlib.AbstractContextManager since Python 3.9.

Answer 5

The Iterator[] version doesn't work when you want to return the contextmanager's reference. For instance, the following code:

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

Will produce an error on the return assert_timing(high=seconds) line:

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

Any legit usage of the function:

with assert_faster_than(1):
    be_quick()

Will result in something like this:

"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__"?

You could fix it like this...

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

But I am going to use the new ContextManager[] object instead and silence out mypy for the decorator:

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 6

Based on PEP-585 the correct annotation type seems to be AbstractContextManager (see https://www.python.org/dev/peps/pep-0585/#implementation). Than you can use the following code:

import contextlib

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

This is the only solution that correctly works together with PyCharm (along with typing.ContextManager, but this one should be deprecated from Python 3.9). It correctly helps you when you use it in with statement (type hints) which is very helpful.

But when I go back to the original question ("How should a context manager be annotated with Python type hints?") it depends. From my point of view the correct one should be the one I mentioned. But this seems not working with mypy (yet). There were some updates regarding this PEP (see https://github.com/python/mypy/issues/7907), but since I'm not much experienced with mypy I might be missing something here.

Answer 7

I had a similar problem when implementing the abstract method:

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


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

Annotating the abstract method with the ContextManager[None] and the implementation with Iterator[None] solves the problem.