Writing a context manager class (Python)

The with statement lesson showed what with promises: the exit step runs whether the block ends normally or with an error. That lesson used managers that the standard library supplies, and it left the machinery for later. You now know enough about classes and dunder methods to build one yourself.

A context manager is any object with two special methods:

__enter__(self) runs when the block starts. Whatever it returns is bound to the name after as.
__exit__(self, exc_type, exc_value, traceback) runs when the block ends, however it ends.

The smallest manager

This class does nothing but report when each method runs and what __exit__ was told:

class Announce:
    def __enter__(self):
        print("enter")
        return self

    def __exit__(self, exc_type, exc_value, traceback):
        if exc_type is None:
            print("exit: no error")
        else:
            print("exit:", exc_type.__name__, "-", exc_value)


with Announce():
    print("working")

print("---")

try:
    with Announce():
        raise ValueError("boom")
except ValueError:
    print("the error carried on")

Output

enter
working
exit: no error
---
enter
exit: ValueError - boom
the error carried on

In the first run no error occurred, so all three arguments of __exit__ were None. In the second, __exit__ received the exception’s type, the exception object, and its traceback. The with statement still ended the block early, and the ValueError carried on to the except outside. A manager that only observes does not interfere.

What `as` receives

The value after as is whatever __enter__ returns, which need not be the manager:

class Greeter:
    def __enter__(self):
        return "hello from __enter__"

    def __exit__(self, exc_type, exc_value, traceback):
        pass


with Greeter() as message:
    print(message)

Output

hello from __enter__

Most managers return self, so the name after as gives access to the manager’s state and methods. Some return a different object, as open() does when it gives you the file.

A manager that cleans up: commit or roll back

Because __exit__ is told whether an error happened, it can behave differently on success and on failure. This one makes a copy of a dictionary on entry. If the block fails, it restores the copy:

class Transaction:
    def __init__(self, data):
        self.data = data

    def __enter__(self):
        self._snapshot = dict(self.data)
        return self.data

    def __exit__(self, exc_type, exc_value, traceback):
        if exc_type is None:
            print("committed")
        else:
            self.data.clear()
            self.data.update(self._snapshot)
            print("rolled back")
        return False


balances = {"a": 100, "b": 0}

with Transaction(balances) as t:
    t["a"] -= 50
    t["b"] += 50
print(balances)

try:
    with Transaction(balances) as t:
        t["a"] -= 80
        raise ValueError("insufficient funds")
except ValueError as error:
    print("error:", error)
print(balances)

Output

committed
{'a': 50, 'b': 50}
rolled back
error: insufficient funds
{'a': 50, 'b': 50}

The first block finished normally, so the changes were kept. In the second, the ValueError arrived with the dictionary already changed, __exit__ restored it, and the error was still reported to the except. This is the same pattern that database libraries use to protect against half-finished work.

__exit__ returned False, so the exception was not swallowed. That choice deserves its own section.

Suppressing exceptions

The return value of __exit__ decides what happens to an exception that left the block:

`__exit__` returns	Result
a false value (`False`, `None`, or nothing)	the exception continues out of the `with`
`True`	the exception is swallowed and the code after the block runs

Falling off the end of __exit__ returns None, so by default errors propagate. If you do want to swallow some errors, be selective. This class ignores the exception types you hand it, as contextlib.suppress from the earlier lesson does:

class Ignore:
    def __init__(self, *exc_types):
        self.exc_types = exc_types

    def __enter__(self):
        return self

    def __exit__(self, exc_type, exc_value, traceback):
        return exc_type is not None and issubclass(exc_type, self.exc_types)


with Ignore(KeyError):
    {}["missing"]
print("after the first block")

try:
    with Ignore(KeyError):
        1 / 0
except ZeroDivisionError:
    print("ZeroDivisionError was not ignored")

Output

after the first block
ZeroDivisionError was not ignored

The KeyError was swallowed and execution resumed after the block. The ZeroDivisionError was not on the list, so __exit__ returned False and the error went on to the outer except. The *exc_types in __init__ collects any number of arguments into a tuple, and issubclass accepts a tuple, so the manager works for several types at once.

The expression exc_type is not None and ... matters. On a normal exit exc_type is None, and issubclass(None, ...) would raise a TypeError.

Common pitfalls

Forgetting return self. If __enter__ returns nothing, the name after as is None:

class Resource:
    def __enter__(self):
        print("opened")  # forgot: return self

    def __exit__(self, exc_type, exc_value, traceback):
        print("closed")

    def read(self):
        return "data"


try:
    with Resource() as r:
        print(r.read())
except AttributeError as error:
    print("AttributeError:", error)

Output

opened
closed
AttributeError: 'NoneType' object has no attribute 'read'

Notice that the manager still opened and closed correctly. Only the as target was wrong, and it failed at the first use.

Returning True carelessly. A manager whose __exit__ always returns True hides every error, including your own bugs:

class Quiet:
    def __enter__(self):
        return self

    def __exit__(self, exc_type, exc_value, traceback):
        return True  # swallows every error, including bugs


with Quiet():
    print(undefined_name)
print("the NameError vanished")

Output

the NameError vanished

The NameError was a genuine bug, and the program went on as if nothing had happened. Return True only for the specific exceptions the manager exists to handle.

Raising from __exit__. If cleanup itself fails, its error replaces the original one, and the cause is harder to see. Keep __exit__ simple, and make cleanup steps that cannot fail where you can.

Reusing a manager instance. Some managers, like the transaction above, keep state on self between __enter__ and __exit__. Entering the same instance again from inside its own block overwrites that state. Create a fresh instance for each with.

ExerciseA timer

Write a class Timer that records how long its block took. __enter__ stores the start time from time.perf_counter() and returns self. __exit__ stores the elapsed seconds in self.elapsed and does not swallow errors.

Run this script, which sleeps for 50 milliseconds inside the block. The elapsed time varies from run to run, so the script prints comparisons, not the number itself. Expected output:

import time

with Timer() as t:
    time.sleep(0.05)
print(t.elapsed >= 0.04)
print(t.elapsed < 5)

Output

True
True

Show solution

import time


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

    def __exit__(self, exc_type, exc_value, traceback):
        self.elapsed = time.perf_counter() - self.start
        return False


with Timer() as t:
    time.sleep(0.05)
print(t.elapsed >= 0.04)
print(t.elapsed < 5)

Output

True
True

time.perf_counter() is a clock meant for measuring durations. __exit__ computes the elapsed time even if the block raised, and then returns False, so the error still propagates to the caller.

ExerciseSection markers

Write a class Section(title) that prints [title] on entry. On a normal exit it prints [/title], and when the block failed it prints [/title failed: ErrorName], using the exception’s class name. It must not swallow errors.

Test it with a block that prints reading, then with a block that calls int("x") inside a try, printing handled afterwards. Expected output:

[load]
reading
[/load]
[parse]
[/parse failed: ValueError]
handled

Show solution

class Section:
    def __init__(self, title):
        self.title = title

    def __enter__(self):
        print(f"[{self.title}]")
        return self

    def __exit__(self, exc_type, exc_value, traceback):
        if exc_type is None:
            print(f"[/{self.title}]")
        else:
            print(f"[/{self.title} failed: {exc_type.__name__}]")
        return False


with Section("load"):
    print("reading")

try:
    with Section("parse"):
        int("x")
except ValueError:
    print("handled")

Output

[load]
reading
[/load]
[parse]
[/parse failed: ValueError]
handled

exc_type is None on success and the exception class on failure, so a single if chooses the message. exc_type.__name__ gives ValueError as text. Returning False lets the ValueError reach the outer except, which is what prints handled.

ExercisePredict the nesting

Two managers are nested, and the inner block raises an error. Write the lines this prints, in order, before running it. Which manager exits first?

class Step:
    def __init__(self, name):
        self.name = name

    def __enter__(self):
        print("enter", self.name)
        return self

    def __exit__(self, exc_type, exc_value, traceback):
        print("exit", self.name)
        return False


try:
    with Step("outer"):
        with Step("inner"):
            print("body")
            raise RuntimeError("stop")
except RuntimeError:
    print("caught")

Show solution

enter outer
enter inner
body
exit inner
exit outer
caught

Managers exit in the reverse order of entry, like nested function calls unwinding. The inner manager exits first, the error then passes outwards through the outer manager, which also exits, and only then does the except clause run.

Summary

A context manager is an object with __enter__ and __exit__. The with statement calls them at the start and end of the block.
The name after as is whatever __enter__ returns. Return self unless you have a reason not to.
__exit__ receives the exception type, value and traceback, all None when no error happened.
A false return from __exit__ lets an exception continue. True swallows it, so use it only for the specific errors the manager is designed to handle.
Nested managers exit in reverse order.

The smallest manager

What as receives

A manager that cleans up: commit or roll back

Suppressing exceptions

Common pitfalls

Summary

What `as` receives