Context managers

What you'll learn

• Explain what with does: enter, run the block, exit, even on error
• Write a context manager as a class with __enter__ and __exit__
• Write one as a generator with @contextlib.contextmanager, which is shorter
• Use the standard library's ready-made managers and combine several in one with

Concept

You have written with open(path) as f: since the files lesson and trusted that the file closes. The with statement works with any object that knows how to set up before a block and clean up after it, no matter how the block ends. Files, locks, database transactions, temporary folders, timers, and "change this setting for a while" are all the same shape: acquire, use, release. A context manager is an object that implements that shape.

The protocol is two methods. __enter__ runs when the block starts and its return value is what as binds; for a file that is the file itself. __exit__ runs when the block ends, and receives three arguments describing the exception if one occurred, or three Nones if not. It does the cleanup, and if it returns True the exception is swallowed; normally you return nothing so errors propagate after cleanup. This is try/finally packaged as a reusable object, which is why it is the tool whenever cleanup must not be forgotten.

Writing a class for every manager is verbose, so contextlib.contextmanager lets a generator function do the job. Everything before yield is the setup, the yielded value is what as binds, and everything after yield is the cleanup. Wrap the yield in try/finally so cleanup runs on errors too.

So the same protocol has two shapes, and choosing between them is mostly a question of size. Most context managers you write should take the generator form; use the class form when the manager needs several methods or holds state that other code inspects after the block.

The standard library ships many. open, threading.Lock, sqlite3 connections, tempfile.TemporaryDirectory, contextlib.suppress(ExcType) to ignore a specific exception, contextlib.redirect_stdout to capture prints, and contextlib.chdir to change the working folder temporarily. Several managers go in one with separated by commas, or in parentheses across lines since 3.10, and contextlib.ExitStack manages a variable number of them when you do not know how many until runtime.

Code

The class form: a timer that reports how long its block took, and reports even when the block raises.

# timer_class.py
import time


class Timer:
    def __init__(self, label):
        self.label = label

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

    def __exit__(self, exc_type, exc, tb):
        self.elapsed = time.perf_counter() - self.start
        status = "failed" if exc_type else "ok"
        print(f"{self.label}: {status}, under {int(self.elapsed * 1000) // 50 * 50 + 50} ms")
        return False


with Timer("sum") as t:
    total = sum(range(200_000))
print(total, t.elapsed >= 0)

try:
    with Timer("broken"):
        raise ValueError("boom")
except ValueError as err:
    print("propagated:", err)

sum: ok, under 50 ms
19999900000 True
broken: failed, under 50 ms
propagated: boom

The generator form, which is the one you will write most: a temporary working folder that is removed afterwards, and a manager that changes a setting for the duration of the block.

# generator_form.py
import contextlib
import shutil
from pathlib import Path

SETTINGS = {"verbose": False}


@contextlib.contextmanager
def scratch_dir(name):
    path = Path(name)
    path.mkdir(exist_ok=True)
    try:
        yield path
    finally:
        shutil.rmtree(path)
        print(f"removed {name}")


@contextlib.contextmanager
def setting(key, value):
    old = SETTINGS[key]
    SETTINGS[key] = value
    try:
        yield
    finally:
        SETTINGS[key] = old


with scratch_dir("tmp_work") as d:
    (d / "a.txt").write_text("x")
    print(sorted(p.name for p in d.iterdir()))
print(Path("tmp_work").exists())

with setting("verbose", True):
    print("inside:", SETTINGS["verbose"])
print("after:", SETTINGS["verbose"])

['a.txt']
removed tmp_work
False
inside: True
after: False

Standard-library managers: suppressing one expected exception, capturing printed output, and a temporary folder, combined in a single with.

# stdlib_managers.py
import contextlib
import io
import tempfile
from pathlib import Path

with contextlib.suppress(FileNotFoundError):
    Path("does-not-exist.txt").unlink()
print("no error from a missing file")

buffer = io.StringIO()
with contextlib.redirect_stdout(buffer):
    print("captured line")
print("captured:", repr(buffer.getvalue()))

with (
    tempfile.TemporaryDirectory() as tmp,
    open(Path(tmp) / "note.txt", "w", encoding="utf-8") as f,
):
    f.write("hello")
    folder = tmp
print(Path(folder).exists())

no error from a missing file
captured: 'captured line\n'
False

ExitStack for a variable number of resources: open every file in a list, and guarantee that all of them close, however many there are.

# exit_stack.py
import contextlib
from pathlib import Path

names = ["a.txt", "b.txt", "c.txt"]
for name in names:
    Path(name).write_text(name.upper(), encoding="utf-8")

with contextlib.ExitStack() as stack:
    files = [stack.enter_context(open(name, encoding="utf-8")) for name in names]
    print([f.read() for f in files])
    print([f.closed for f in files])
print([f.closed for f in files])

['A.TXT', 'B.TXT', 'C.TXT']
[False, False, False]
[True, True, True]

Try it yourself

Create managers.py with a @contextmanager called logged(label) that prints enter <label> before the block and exit <label> after, even if the block raises, and a class-based Transaction that collects operations in a list and prints "committed" on a clean exit or "rolled back" when an exception occurs. Use both together in one with statement, once succeeding and once with a deliberate error caught outside.

Common mistakes

• Forgetting try/finally around yield. Without it, an exception in the block skips the cleanup after the yield.
• Returning True from __exit__ by accident. That swallows every exception; return False or nothing unless suppression is the point.
• Writing a class when a generator would do. Three lines with @contextmanager beat a ten-line class for most setup-and-cleanup pairs.

What's next

• with calls __enter__, runs the block, and always calls __exit__, which receives the exception details if any.
• @contextlib.contextmanager turns a generator into a manager: setup, yield, cleanup in finally.
• The standard library provides suppress, redirect_stdout, TemporaryDirectory, and ExitStack; several managers fit in one with.

You can guarantee cleanup for any resource. Next you will make your objects behave like built-ins, printing, comparing, hashing, indexing, and adding, with the data model's dunder methods.