Create your own Python context manager

Organic and contextually sourced

Jul 15, 2023

Summary

Creating a custom context manager takes little time and increases code quality, even if you use it only once.

It's as quick as writing this:

class YourContextManager:
    def __enter__(self):
        print("Setup up")
    def __exit__(self, exc_type, exc_value, traceback):
        print("Tear down")

And if this is too annoying to type, you can even do:

from contextlib import contextmanager
@contextmanager
def your_context_manager():
    print("Setup up")
    yield
    print("Tear down")

But, why?

Last article, we saw how to use context managers and today, we are going to create our own custom ones. The reason to do so is purely to help with code quality. You will encounter situations where you need setup and tear down for which no context manager exists. To me, context managers are an exception to DRY. By this I mean I would often code one even if I plan to use it a single time.

This is because I don't see context managers only as a good tool for reusable code, but also something:

That makes the code easier to read and test;
Force you to delimitate the scope of things.
With their sheer presence, that prevents future code editions from introducing certain kinds of bugs.
Encourage good style, even for developers new to the code base.

The class-based syntax

Python uses dunder methods (for "double underscore methods") to let the programmer create a context manager:

>>> import time
... class House:
...
...     def __enter__(self):
...         # This is called when using "with"
...         print("Knock, knock!")
...
...     def __exit__(self, exc_type, exc_value, traceback):
...         # This is called at the end of the "with" block
...         time.sleep(10)
...         print("...Python")
...
>>> with House():
...     print("Who's there?")
...
Knock, knock!
Who's there?
...Python

Now there are a few things to understand.

First, __enter__ is executed immediately after you encounter the with. If you return a value from it, the user can bind it using as:

>>> import time
... class House:
...     def __init__(self, wait=10):
...         self.wait = wait
...
...     def __enter__(self):
...         print("Knock, knock!")
...         return self.wait
...
...     def __exit__(self, exc_type, exc_value, traceback):
...         time.sleep(self.wait)
...         print("...Python")
...
>>> with House(2) as waiting:
...     print("Who's there?")
...     print(f'(waiting {waiting} seconds)')
...
Knock, knock!
Who's there?
(waiting 2 seconds)
...Python

As I'm reading this, I can see how painful it is that substack still doesn't have syntax highlighting.

Second, __exit__ parameters always expect three parameters:

exc_type, exc_value and traceback.

The names don't matter, and in fact, if you don't remember them, you can just *args your way out of it. But good names are important, so do it sparingly. Those variables will either all be set to None (if the with block executed without error), or contains info about any error that would have interrupted the with block:

exc_type would contain the class of the exception. E.G: ValueError.
exc_value would contain the instance of the exception. E.G: ValueError("Yeah, you really should not do that").
traceback will contain all the data you need to reconstitute the stack trace.

So __exit__ is always called, whether you had an error or not:

>>> import time
... class House:
...
...     def __enter__(self):
...         print("Knock, knock!")
...
...     def __exit__(self, exc_type, exc_value, traceback):
...         time.sleep(0)
...         print("...Rust")
...         if exc_type is not None:
...            print(f'Woops! We had an error: {exc_value}')
...
>>> with House():
...     print("Who's there?")
...     1/0
...
Knock, knock!
Who's there?
...Rust
Woops! We had an error: division by zero
Traceback (most recent call last):
  Cell In[8], line 15
    1/0
ZeroDivisionError: division by zero

Despite the exception, __exit__ has not only be running successfully, but it had the opportunity to inspect the error.

In fact, if you return True from __exit__, you can even decide to stop the exception there:

>>> import time
... class House:
...
...     def __enter__(self):
...         print("Knock, knock!")
...
...     def __exit__(self, exc_type, exc_value, traceback):
...         time.sleep(0)
...         print("...Rust")
...         if exc_type is not None:
...            print(f'Woops! We had an error: {exc_value}')
...            print('But life goes on...')
...            return True
...
>>> with House():
...     print("Who's there?")
...     1/0
...
Knock, knock!
Who's there?
...Rust
Woops! We had an error: division by zero
But life goes on...

See? No stack trace.

The generator-based syntax

A lot of context managers are glorified try/except syntactic sugar, and writing a whole class for that is overkill. For this reason, contextlib provides the contextmanager decorator to make writing them easier.

Now, this feature uses decorators and generators to create a context manager, so it's a lot to unpack. If you don't understand this part, it's OK, you can skim-read it. I will write about decorators and generators in other articles.

Let's rewrite the House context manager with it. This:

... class House:
...     def __init__(self, wait=10):
...         self.wait = wait
...
...     def __enter__(self):
...         print("Knock, knock!")
...         return self.wait
...
...     def __exit__(self, exc_type, exc_value, traceback):
...         time.sleep(self.wait)
...         print("...Python")

Becomes:

>>> import time
... from contextlib import contextmanager
... @contextmanager
... def house(wait=5):
...     # this is like __enter__
...     print("Knock, knock!")
...     # This is like return self.wait
...     # The yield is mandatory, but the variable "wait"
...     # is optional.
...     yield wait
...     # This is like __exit__
...     time.sleep(wait)
...     print("...Pypy")
...
>>> with house():
...     print("Who's there?")
...
Knock, knock!
Who's there?
...Pypy

But wait, how do you handle an exception in there, since we don't have access to exc_type, exc_value, or traceback? Well, with a regular try/except:

>>> import time
... from contextlib import contextmanager
... @contextmanager
... def house(wait=5):
...     print("Knock, knock!")
...     try:
...         yield wait
...     # A rare case where using Exception is OK
...     except Exception as ex:
...         print(f'Woops! We had an error: {ex}')
...         # re-raise here if you want the exception to propagate
...     time.sleep(wait)
...     print("...Pypy")
...
>>> with house():
...     print("Who's there?")
...     1/0
...
Knock, knock!
Who's there?
Woops! We had an error: division by zero
...Pypy

More flexible context managers

With those you have enough to cover most of your context manager needs, but we can do better.

One thing to consider is to add methods to manually setup and tear down your context manager, and call them in __enter__ and __exit__:

>>> import time
... class House:
...
...
...     def __enter__(self):
...         self.enter()
...
...     def __exit__(self, exc_type, exc_value, traceback):
...         self.exit()
...
...     def enter(self):
...         print("Knock, knock!")
...
...     def exit(self):
...         time.sleep(10)
...         print("...Python")

Indeed, this will let advanced users control their workflow manually if they need to, by calling .enter() and .exit() manually, instead of using with. This is also useful in the case of inversion of control (typically pub/sub) and works well with some functional programming patterns.

It's also about making your API easy to understand and discover. If people see in their completion list a method .connect() and .disconnect(), they will know what they do, and then can deduct what __enter__() and __exit__() do. If you only see __enter__() and __exit__(), this doesn't say much about what they do and reading the doc becomes necessary.

Another thing you can do is to also provide __aenter__ and __aexit__. This is not a typo, those are the asynchronous equivalent to __enter__ and __exit__. They only make sense if your object is compatible with something like asyncio, but if it is, then your users will be able to use asyncio with instead of with, which will not block.

Since we are at it, you may use contextlib.ContextDecorator to create a context manager that can also be used as a decorator. After all, both are made to execute some code before and after something. Context managers will execute something before and after a block, while decorators will execute something before and after a function call. So they have similarities:

from contextlib import ContextDecorator

class make_sure(ContextDecorator):
    def __enter__(self):
        print("IT HAS BEGUN!")
        return self

    def __exit__(self, *exc):
        print("IT IS OVER!")

This can be used as you expect:

>>> with make_sure():
...     print('<insert some quest here>')
...
IT HAS BEGUN!
<insert some quest here>
IT IS OVER!

And this can also be used as a decorator:

>>> @make_sure() # don't forget the parenthesis here
... def perform_quest():
...     print('<insert some quest here>')
...
>>> perform_quest()
IT HAS BEGUN!
<insert some quest here>
IT IS OVER!

Finally, since context managers are just regular Python object, you can return one from a call to turn any method into a process with setup and tear down:

class House:
    def __enter__(self):
        print('Knock, knock!')

    def __exit__(*args):
        print('Good bye!')

class Salesman:
    def visit(self):
        return House()

bob = Salesman()
with bob.visit():
    print('Every American should have a vacuum cleaner in their basement')
Knock, knock!
Every American should have a vacuum cleaner in their basement
Good bye!

Bite code!

Discussion about this post