Context managers

Why this matters

Many resources must be released after use: files must be closed, locks released, database transactions committed or rolled back, temporary changes undone. The hard part is doing it on every way out, including when an exception flies through. try/finally works but clutters code, and it is easy to forget. The with statement packages the setup and the guaranteed cleanup into a reusable object, a context manager. You have used with open(...) since the files lesson; now you will see the protocol behind it and write your own, both as classes and with a generator shortcut.

What with does

with open("notes.txt", "w") as f:
    f.write("hello\n")
# f is closed here, whatever happened inside

The with statement is roughly this code:

mgr = open("notes.txt", "w")
f = mgr.__enter__()               # setup; the result is bound by "as"
try:
    f.write("hello\n")            # the body
except BaseException as e:
    if not mgr.__exit__(type(e), e, e.__traceback__):
        raise                     # not suppressed: continue propagating
else:
    mgr.__exit__(None, None, None)

So the protocol is two methods:

• __enter__(self) runs at the start; its return value is what as binds (it does not have to be self).
• __exit__(self, exc_type, exc, tb) runs at the end, always. If the body finished normally the three arguments are None. If an exception was raised, they describe it. If __exit__ returns a true value, the exception is suppressed; otherwise it continues after cleanup.

__exit__ also runs when you leave the block with return, break or continue.

A class-based context manager

class Tag:
    def __init__(self, name):
        self.name = name
    def __enter__(self):
        print(f"<{self.name}>")
        return self
    def __exit__(self, exc_type, exc, tb):
        print(f"</{self.name}>")
        return False

with Tag("html"):
    with Tag("body"):
        print("hi")

<html>
<body>
hi
</body>
</html>

Nested with blocks unwind in reverse order, like a stack. Several managers can go on one line: with Tag("a"), Tag("b"):, or, in parentheses, across lines.

Now with an exception:

try:
    with Tag("p"):
        print(1 / 0)
except ZeroDivisionError as e:
    print("caught:", e)

<p>
</p>
caught: division by zero

The closing tag still printed: __exit__ ran, returned False, and the exception continued to the outer try.

Suppressing exceptions

Returning True from __exit__ swallows the exception, and execution continues after the with block. This is powerful and dangerous; do it only for specific exception types. The standard library has a ready-made one:

import contextlib, os

with contextlib.suppress(FileNotFoundError):
    os.remove("maybe-missing.tmp")
print("carried on")

carried on

That replaces a four-line try/except FileNotFoundError: pass.

The generator shortcut: contextlib.contextmanager

Writing a class for every context manager is verbose. contextlib.contextmanager turns a generator function with a single yield into one:

from contextlib import contextmanager

@contextmanager
def tag(name):
    print(f"<{name}>")         # __enter__ part
    try:
        yield name             # the value bound by "as"; the body runs here
    finally:
        print(f"</{name}>")    # __exit__ part

with tag("b") as t:
    print("bold", t)

<b>
bold b
</b>

How it works: the decorator wraps your generator in an object whose __enter__ calls next() to run up to the yield, and whose __exit__ resumes the generator. If the body raised, __exit__ throws the exception into the generator at the yield, which is why you need try/finally (or try/except) around the yield to guarantee cleanup. Without it, an exception in the body would skip your cleanup code. To suppress an exception, catch it around the yield and do not re-raise.

A practical example, timing a block (output values vary, so we just show the shape):

import time
from contextlib import contextmanager

@contextmanager
def timer(label):
    start = time.perf_counter()
    try:
        yield
    finally:
        print(f"{label}: {time.perf_counter() - start:.3f}s")

with timer("sum"):
    total = sum(range(10_000_000))

sum: 0.142s

Where you will meet context managers

• open(...): closes the file.
• threading.Lock(): with lock: acquires and releases (concurrency module).
• sqlite3 connections: with conn: commits on success, rolls back on exception.
• tempfile.TemporaryDirectory(): deletes the directory afterward.
• contextlib.redirect_stdout(buf): temporarily sends print output elsewhere, handy in tests.
• unittest's with self.assertRaises(ValueError): checks that the block raises.

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

'captured\n'

ExitStack: a dynamic number of managers

When you do not know until run time how many resources you need, for example opening every file in a list, contextlib.ExitStack collects managers and exits them all, in reverse order, when its own block ends:

from contextlib import ExitStack

names = ["a.txt", "b.txt", "c.txt"]
with ExitStack() as stack:
    files = [stack.enter_context(open(n, "w")) for n in names]
    for f in files:
        f.write("data\n")
# all three files are closed here, even if one open() failed halfway

stack.callback(func, *args) also registers arbitrary cleanup functions.

Common mistakes

• No try/finally around yield in a @contextmanager function: cleanup is skipped when the body raises.
• Accidentally suppressing everything: __exit__ returning True unconditionally hides every bug inside the block.
• Yielding twice, or never: RuntimeError: generator didn't stop or generator didn't yield. A @contextmanager generator must yield exactly once.
• Using the resource after the block: f.read() after with open(...) as f: raises ValueError: I/O operation on closed file.
• Forgetting that as binds __enter__'s return value: if __enter__ returns nothing, as x gives None.

Recap

• with mgr as x: calls mgr.__enter__() (bound to x), runs the body, then always calls mgr.__exit__(exc_type, exc, tb).
• A true return value from __exit__ suppresses the exception; prefer letting exceptions propagate.
• @contextlib.contextmanager turns a one-yield generator into a context manager; protect the yield with try/finally.
• contextlib.suppress, redirect_stdout and ExitStack cover common needs.

Going further: read the contextlib documentation and PEP 343, "The with Statement", which shows the exact translation of with into try/finally.