Python Type Hints - How to Type a Context Manager2021-07-04
Python’s context manager protocol has only two methods, with straightforward types. But when it comes to adding accurate type hints to a context manager, we still need to combine several typing features. Let’s look at how we can do this for the two different ways of making a context manager.
The easiest way to create a context manager is using the
@contextmanager decorator from
When adding type hints, this remains the easiest way, as we only need to type the underlying generator.
collections.abc.Generator is only supported on Python 3.9; on older versions we need to import
Generator to specify the three types our generator uses - the yield type, send type, and return type.
In this case we do not yield a value, so we set the yield type to
@contextmanager never sends anything into our generator, and ignores our return value, so we should always use
None for the second and third values.
For context managers that return values, we would swap the first
None for the value’s type, for example:
Note: the documentation for
typing.Generator notes simple generators can use e.g.
This is currently accepted by the
@contextmanager type hints in typeshed, but there’s an open issue showing how this can lead to bugs.
Therefore it’s best to stick to
Thanks to Tom Grainger for pointing this out on Twitter, and Anthony Sottile for reporting the issue.
Class-based context managers
We can create more complicated context managers as classes, and this is where the type hints need a bit more work. The simplest definition looks like this:
For context managers that return values, we would swap
__enter__’s return type from
None to the value’s type.
For context managers that suppress exceptions, we would change
__exit__’s return type to
__exit__ method’s type hints say:
exc_typeis a class that inherits from
exc_valis an instance of
BaseException(or a subclass), or
exc_tbis a traceback, or
These are all true, but the type hints don’t represent the correlation between the variables: they’re either all set, or all
__exit__ can never be called with only some values not
If we only care about using this context manager with the
with statement, we can handle this correlation as needed with type narrowing inside our
__exit__ method’s body.
But if we care about users calling
__exit__ directly, we can reach for
Let’s look at these two techniques in turn.
Type narrowing in
To handle the correlation inside we add some type narrowing with
If we wanted to handle exceptions, we could use a body like this:
Mypy can use the
if statement to infer the type of
exc_type in both blocks: it’s
type[BaseException] in the
if block, and
None in the
But because Mypy doesn’t know about the correlation between the variables, it can’t narrow the types of
We can tell help Mypy narrow the types with
Mypy can read the
asserts and determine the variables’ types in the following lines.
The above example contains the complete set of
assert statements, but we don’t always need to be so exhaustive.
If a block doesn’t use a variable, we don’t need to narrow its type there.
For example, many context managers only use the exception value, which we can do without mentioning the other variables:
If in doubt about the type narrowing you need, debug with
To make the correlation between
__exit__’s arguments visible to callers, we need to use
@typing.overload to list the accepted forms.
This requires a couple extra stub functions:
The first two
@overload-decorated functions declare the allowed types for callers.
We spell out the two cases: either all the arguments are
None, or all the arguments are set.
__exit__ function is the implementation, and here we need to combine the overloaded types.
Note that, since we have to use unions, inside the body we still need to use type narrowing as above.
(Mypy can’t propagate the
@overload information into the body at current.)
Now if callers try to pass an incomplete set of arguments, they will get a type error. For example, if we wrote a call like this:
Then Mypy would complain like so:
I hope this has managed to give you some context,
🎉 My book Speed Up Your Django Tests is now up to date for Django 3.2. 🎉
Buy now on Gumroad
One summary email a week, no spam, I pinky promise.
- Python Type Hints - How to Specify a Class Rather Than an Instance Thereof
- Python Type Hints - How to Narrow Types with isinstance(), assert, and Literal
- Python Type Hints - How to Use @overload
Tags: mypy, python
© 2021 All rights reserved.