Python type hints and context managers

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

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

from contextlib import contextmanagerfrom typing import ContextManager@contextmanagerdef session() -> ContextManager[Session]:    yield Session(...)

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

from typing import Iteratordef assert_faster_than(seconds: float) -> Iterator[None]:    return assert_timing(high=seconds)@contextmanagerdef 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 ContextManagerdef assert_faster_than(seconds: float) -> ContextManager[None]:    return assert_timing(high=seconds)@contextmanager  # type: ignoredef assert_timing(low: float = 0, high: float = None) -> ContextManager[None]:    ...