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

Answer 1

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