Python type hints and context managers

Solution 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.

Solution 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

Solution 3:

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

Solution 4:

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.