Python type hints and context managers

python mypy
Peter picture Peter · Apr 9, 2018 · Viewed 12.2k times · Source

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

Michael0x2a picture Michael0x2a · Apr 9, 2018

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.

Related questions

How can I specify the function type in my type hints?

I want to use type hints in my current Python 3.5 project. My function should receive a function as parameter. How can I specify the type function in my type hints? import typing def my_function(name:typing.AnyStr, func: typing.…

Read More
How can mypy ignore a single line in a source file?

I'm using mypy in my python project for type checking. I'm also using PyYAML for reading and writing the project configuration files. Unfortunately, when using the recommended import mechanism from the PyYAML documentation this generates a spurious error in a …

Read More
Specifying a type to be a List of numbers (ints and/or floats)?

How do I specific a function can take a list of numbers which can be ints or floats? I tried making a new type using Union like so: num = Union[int, float] def quick_sort(arr: List[num]) -> …

Read More