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

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

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)

Solution 4

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 5

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.

View more solutions

Comments

[internet_user]

It's a generator, and it yields, sends, and returns None, so it's a Generator[None, None, None]. It doesn't matter if you use it for a context manager.

[Onilol]

If you have any idea on what this specific context manager will be used for, you can annotate for the expected types, else you'd be pretty much accepting anything (even None)

[Peter]

In my specific case I just want to use the context manager for logging (timing) so the yield, send and return values really are None.

[shmee]

Looking into a potential dupe, I came across this answer. It seems like the return type for a generator used in a context manager should reflect what the context manager returns, i.e. ContextManager[_T]. With that, the static checker in my IDE was able to successfully infer the type of the context variable, while it did not work with Iterator. Can you check? I'd like to flag the other question as a dupe, but as it stands, this answer does not solve the problem in the other Q.

[CMCDragonkai]

This doesn't seem to work for me. Mypy says error: The return type of a generator function should be "Generator" or one of its supertypes and error: Argument 1 to "contextmanager" has incompatible type "Callable[[Abc, Any, Any], ContextManager[Any]]"; expected "Callable[..., Iterator[<nothing>]]"

[Marius Gedminas]

You want the type signatures of assert_faster_than and assert_timing to look the same, but you're applying @contextmanager to only one of them. I think the right thing to do is to declare assert_faster_than(...) -> ContextManager[None], but assert_timing(..) -> Iterator[None].

[kolypto]

I guess mypy is too strict :D I don't have a better annotation at the moment

[Robino]

Type hinting now works for me thanks to this. PyCharm (2020.1.2 Community Edition) and python 3.8.

[eric.frederich]

Thanks, this helped with PyCharm but not mypy. Perhaps a single solution does not yet exist to make both tools happy

[levsa]

This gives en error with Python 3.7.9 (when running the code): TypeError: 'ABCMeta' object is not subscriptable

[Nerxis]

@levsa: This PEP is meant for Python 3.9 and newer, if you want to try this for older Python versions (from 3.7) you have to use from __future__ import annotations to be forward-compatible.

[Neil G]

@kolypto No, it's not that mypy is too strict. PyCharm is simply wrong. You should be annotating it as Generator, and the decorator will take that Generator and return a ContextManager.

[Dustin Wyatt]

@shmee I'm not sure I agree that "the return type for a generator used in a context manager should reflect what the context manager returns". The function returns what it returns, and I usually think of the decorator as modifying the function...so if you want to know what the decorated function returns you need to look at the type annotations for the decorator.