Creating a context manager in Python

How can you create your own context manager in Python?

What is a context manager?

A context manager is an object that can be used in a with block to sandwich some code between an entrance action and an exit action.

File objects can be used as context managers to automatically close the file when we're done working with it:

>>> with open("example.txt", "w") as file:
...     file.write("Hello, world!")
...
13
>>> file.closed
True

Context managers need a __enter__ method and a __exit__ method, and the __exit__ method should accept three positional arguments:

class Example:
    def __enter__(self):
        print("enter")

    def __exit__(self, exc_type, exc_val, exc_tb):
        print("exit")

This context manager just prints enter when the with block is entered and exit when the with block is exited:

>>> with Example():
...     print("Yay Python!")
...
enter
Yay Python!
exit

Of course, this is a somewhat silly context manager. Let's look at a context manager that actually does something a little bit useful.

A useful context manager

This context manager temporarily changes the value of an environment variable:

import os

class set_env_var:
    def __init__(self, var_name, new_value):
        self.var_name = var_name
        self.new_value = new_value

    def __enter__(self):
        self.original_value = os.environ.get(self.var_name)
        os.environ[self.var_name] = self.new_value

    def __exit__(self, exc_type, exc_val, exc_tb):
        if self.original_value is None:
            del os.environ[self.var_name]
        else:
            os.environ[self.var_name] = self.original_value

The USER environment variable on my machine currently has the value of Trey:

>>> print("USER env var is", os.environ["USER"])
USER env var is trey

If we use this context manager, within its with block, the USER environment variable will have a different value:

>>> with set_env_var("USER", "akin"):
...     print("USER env var is", os.environ["USER"])
...
USER env var is akin

But after the with block exits, the value of that environment variable resets back to its original value:

>>> print("USER env var is", os.environ["USER"])
USER env var is trey

This is all thanks to our context manager's __enter__ method and a __exit__ method, which run when our context manager's with block is entered and exited.

What about that as keyword?

You'll sometimes see context managers used with an as keyword (note the as result below):

>>> with set_env_var("USER", "akin") as result:
...     print("USER env var is", os.environ["USER"])
...     print("Result from __enter__ method:", result)
...

The as keyword will point a given variable name to the return value from the __enter__ method:

In our case, we always get None as the value of our result variable:

>>> with set_env_var("USER", "akin") as result:
...     print("USER env var is", os.environ["USER"])
...     print("Result from __enter__ method:", result)
...
USER env var is akin
Result from __enter__ method: None

This is because our __enter__ method doesn't return anything, so it implicitly returns the default function return value of None.

The return value of __enter__

Let's look at a context manager that does return something from __enter__.

Here we have a program called timer.py:

import time

class Timer:
    def __enter__(self):
        self.start = time.perf_counter()
        return self

    def __exit__(self, exc_type, exc_val, exc_tb):
        self.stop = time.perf_counter()
        self.elapsed = self.stop - self.start

This context manager will time how long it took to run a particular block of code (the block of code in our with block).

We can use this context manager by making a Timer object, using with to run a block of code, and then checking the elapsed attribute on our Timer object:

>>> t = Timer()
>>> with t:
...     result = sum(range(10_000_000))
...
>>> t.elapsed
0.28711878502508625

But there's actually an even shorter way to use this context manager.

We can make the Timer object and assign it to a variable, all on one line of code, using our with block and the as keyword:

>>> with Timer() as t:
...     result = sum(range(10_000_000))
...
>>> t.elapsed
0.3115791230229661

This works because our context manager's __enter__ method returns self:

    def __enter__(self):
        self.start = time.perf_counter()
        return self

So it's returning the actual context manager object to us and that's what gets assigned to the t variable in our with block.

Since many context managers keep track of some useful state on their own object, it's very common to see a context manager's __enter__ method return self.

The arguments passed to __exit__

What about that __exit__ method?

    def __exit__(self, exc_type, exc_val, exc_tb):
        self.stop = time.perf_counter()
        self.elapsed = self.stop - self.start

What are those three arguments that it accepts? And does its return value matter?

If an exception occurs within a with block, these three arguments passed to the context manager's __exit__ method will be:

1. the exception class
2. the exception object
3. a traceback object for the exception

But if no exception occurs, those three arguments will all be None.

Here's a context manager that uses all three of those arguments:

import logging

class LogException:
    def __init__(self, logger, level=logging.ERROR, suppress=False):
        self.logger, self.level, self.suppress = logger, level, suppress

    def __enter__(self):
        return self

    def __exit__(self, exc_type, exc_val, exc_tb):
        if exc_type is not None:
            info = (exc_type, exc_val, exc_tb)
            self.logger.log(self.level, "Exception occurred", exc_info=info)
            return self.suppress
        return False

This context manager logs exceptions as they occur (using Python's logging module).

So we can use this LogException context manager like this:

import logging
from log_exception import LogException

logging.basicConfig(level=logging.DEBUG)
logger = logging.getLogger("example")

with LogException(logger):
    result = 1 / 0  # This will cause a ZeroDivisionError
print("That's the end of our program")

When an exception occurs in our code, we'll see the exception logged to our console:

$ python3 log_example.py
ERROR:example:Exception occurred
Traceback (most recent call last):
  File "/home/trey/_/log_example.py", line 8, in <module>
    result = 1 / 0  # This will cause a ZeroDivisionError
             ~~^~~
ZeroDivisionError: division by zero
Traceback (most recent call last):
  File "/home/trey/_/log_example.py", line 8, in <module>
    result = 1 / 0  # This will cause a ZeroDivisionError
             ~~^~~
ZeroDivisionError: division by zero

We see ERROR, the name of our logger (example), Exception occurred, and then the traceback.

In this example, we also see a second traceback, which was printed by Python when our program crashed.

Because our program exited, it didn't actually print out the last line in our program (That's the end of our program).

The return value of __exit__

If we had passed suppress=True to our context manager, we'll see something different happen:

import logging
from log_exception import LogException

logging.basicConfig(level=logging.DEBUG)
logger = logging.getLogger("example")

with LogException(logger, suppress=True):
    result = 1 / 0  # This will cause a ZeroDivisionError
print("That's the end of our program")

Now when we run our program, the exception is logged, but then our program continues onward after the with block:

$ python3 log_example.py
ERROR:example:Exception occurred
Traceback (most recent call last):
  File "/home/trey/_/_/log_example.py", line 8, in <module>
    result = 1 / 0  # This will cause a ZeroDivisionError
             ~~^~~
ZeroDivisionError: division by zero
That's the end of our program

We can see That's the end of our program actually prints out here!

What's going on?

So this suppress argument, it's used by our context manager to suppress an exception:

import logging

class LogException:
    ...

    def __exit__(self, exc_type, exc_val, exc_tb):
        if exc_type is not None:
            info = (exc_type, exc_val, exc_tb)
            self.logger.log(self.level, "Exception occurred", exc_info=info)
            return self.suppress
        return False

If the __exit__ method returns something true or truthy, whatever exception was being raised will actually be suppressed.

By default, __exit__ returns None, just as every function does by default. If we return None, which is falsey, or False, or anything that's falsey, __exit__ won't do anything different from its default, which is to just continue raising that exception.

But if True or a truthy value is returned, the exception will be suppressed.

Aside: what about contextmanager?

Have you ever seen a generator function that somehow made a context manager?

Python's contextlib module includes a decorator which allows for creating context managers using a function syntax (instead of using the typical class syntax we saw above):

from contextlib import contextmanager
import os


@contextmanager
def set_env_var(var_name, new_value):
    original_value = os.environ.get(var_name)
    os.environ[var_name] = new_value
    try:
        yield
    finally:
        if original_value is None:
            del os.environ[var_name]
        else:
            os.environ[var_name] = original_value

Interestingly, this fancy decorator still involves __enter__ and __exit__ under the hood: it's just a very clever helper for creating an object that has those methods.

This contextmanager decorator can sometimes be very handy, though it does have limitations! See the contextmanager decorator for more.

Make context managers with __enter__ & __exit__

Context managers are objects that work in a with block.

You can make a context manager by creating an object that has a __enter__ method and a __exit__ method.

Python also includes a fancy decorator for creating context managers with a function syntax, which I'll cover in a future screencast.