Context Managers and the with Statement

The problem it solves

# Fragile: close() is skipped if read() raises
handle = open("notes.txt", encoding="utf-8")
content = handle.read()
handle.close()

# Correct, and verbose
handle = open("notes.txt", encoding="utf-8")
try:
    content = handle.read()
finally:
    handle.close()

# The same guarantee, built in
with open("notes.txt", encoding="utf-8") as handle:
    content = handle.read()

Anything acquired must be released: files, locks, connections, temporary state. A context manager attaches the release to the block rather than to your memory.

The protocol

class Managed:
    def __enter__(self):
        print("enter")
        return "the value bound by as"

    def __exit__(self, exc_type, exc_value, traceback):
        print("exit")
        return False          # False: do not suppress any exception


with Managed() as value:
    print("inside:", value)

enter
inside: the value bound by as
exit

The exit arguments

class Reporter:
    def __enter__(self):
        return self

    def __exit__(self, exc_type, exc_value, traceback):
        if exc_type is None:
            print("left cleanly")
        else:
            print(f"left with {exc_type.__name__}: {exc_value}")
        return False          # let the exception continue


with Reporter():
    print("all fine")

print("---")

try:
    with Reporter():
        raise ValueError("something broke")
except ValueError:
    print("caught outside")

When the block ends normally, all three arguments are None. When it ends with an exception, they describe it - and __exit__ still runs, which is the whole point.

Suppressing an exception

class Suppress:
    def __init__(self, *exceptions):
        self.exceptions = exceptions

    def __enter__(self):
        return self

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


with Suppress(ZeroDivisionError):
    print(1 / 0)
    print("never reached")

print("execution continues")

from contextlib import suppress

with suppress(FileNotFoundError):
    import os
    os.remove("might-not-exist.txt")

print("carried on")

Returning True from __exit__ swallows the exception silently. Do it only when that is genuinely the manager's purpose, as with contextlib.suppress. Accidentally returning a truthy value hides real failures.

Writing one with contextlib

from contextlib import contextmanager


@contextmanager
def managed():
    print("setup")                 # everything before yield is __enter__
    try:
        yield "the value"          # what `as` binds
    finally:
        print("cleanup")           # everything after is __exit__


with managed() as value:
    print("inside:", value)

print("---")

try:
    with managed():
        raise ValueError("boom")
except ValueError:
    print("cleanup still ran")

This is the form to reach for. One generator replaces a class with two methods, and the try ... finally makes the cleanup guarantee explicit.

Practical context managers

Timing a block

import time
from contextlib import contextmanager


@contextmanager
def timed(label="block"):
    start = time.perf_counter()
    try:
        yield
    finally:
        print(f"{label} took {time.perf_counter() - start:.4f}s")


with timed("summing"):
    total = sum(range(1_000_000))

Temporarily changing state

import os
from contextlib import contextmanager


@contextmanager
def working_directory(path):
    previous = os.getcwd()
    os.chdir(path)
    try:
        yield path
    finally:
        os.chdir(previous)         # restored even if the block raises


@contextmanager
def environment(**overrides):
    saved = {k: os.environ.get(k) for k in overrides}
    os.environ.update({k: str(v) for k, v in overrides.items()})
    try:
        yield
    finally:
        for key, value in saved.items():
            if value is None:
                os.environ.pop(key, None)
            else:
                os.environ[key] = value


with environment(DEBUG="1", MODE="test"):
    print(os.environ["DEBUG"], os.environ["MODE"])

print(os.environ.get("DEBUG"))     # None - restored

Save, change, restore in finally. This shape covers settings, environment variables, working directories, and anything else that is global and must be put back.

Transactions

from contextlib import contextmanager


class Store:
    def __init__(self):
        self.data = {}

    @contextmanager
    def transaction(self):
        snapshot = dict(self.data)
        try:
            yield self
        except Exception:
            self.data = snapshot       # roll back
            print("rolled back")
            raise
        else:
            print("committed")


store = Store()

with store.transaction():
    store.data["a"] = 1
    store.data["b"] = 2

print(store.data)

try:
    with store.transaction():
        store.data["c"] = 3
        raise ValueError("failure halfway")
except ValueError:
    pass

print(store.data)          # c was rolled back

Indented output

from contextlib import contextmanager

INDENT = 0


@contextmanager
def section(title):
    global INDENT
    print("  " * INDENT + title)
    INDENT += 1
    try:
        yield
    finally:
        INDENT -= 1


def log(message):
    print("  " * INDENT + message)


with section("Build"):
    log("compiling")
    with section("Tests"):
        log("unit")
        log("integration")
    log("packaging")

contextlib helpers

from contextlib import suppress, redirect_stdout, closing, ExitStack, nullcontext
import io

# Ignore specific exceptions
with suppress(KeyError):
    {}["missing"]

# Capture printed output
buffer = io.StringIO()
with redirect_stdout(buffer):
    print("captured")
print("got:", buffer.getvalue().strip())

# Optional context manager
def process(path=None):
    manager = open(path, encoding="utf-8") if path else nullcontext()
    with manager as handle:
        return handle.read() if handle else "no file given"


print(process())

from contextlib import ExitStack

paths = ["a.txt", "b.txt", "c.txt"]

with ExitStack() as stack:
    handles = [stack.enter_context(open(p, encoding="utf-8")) for p in paths]
    for handle in handles:
        print(handle.name, len(handle.read()))
# every file is closed, in reverse order

ExitStack manages a number of context managers that is not known until runtime, and closes them all even if one of them fails.

Reusable and reentrant managers

from contextlib import contextmanager


@contextmanager
def once():
    yield "value"


manager = once()

with manager:
    pass

# with manager:      # RuntimeError: generator didn't yield
#     pass

A @contextmanager generator is single use. Call the function again to get a fresh one. A class based manager can be written to be reusable, and threading.RLock is an example of a reentrant one.

Class or generator?

A worked example

import os
import tempfile
from contextlib import contextmanager
from pathlib import Path


@contextmanager
def atomic_write(path, encoding="utf-8"):
    """Write to a temporary file and replace the target only on success."""
    path = Path(path)
    handle = tempfile.NamedTemporaryFile(
        mode="w",
        encoding=encoding,
        dir=path.parent,
        delete=False,
        suffix=".tmp",
    )
    temporary = Path(handle.name)
    try:
        with handle:
            yield handle
        os.replace(temporary, path)          # atomic
    except Exception:
        temporary.unlink(missing_ok=True)    # leave the original untouched
        raise


target = Path("settings.txt")
target.write_text("original\n", encoding="utf-8")

try:
    with atomic_write(target) as handle:
        handle.write("new content\n")
        raise ValueError("failure halfway")
except ValueError:
    pass

print(target.read_text(encoding="utf-8"))    # original - nothing was lost

with atomic_write(target) as handle:
    handle.write("new content\n")

print(target.read_text(encoding="utf-8"))    # new content

Common mistakes

• Returning a truthy value from __exit__ by accident, swallowing exceptions.
• Omitting the try ... finally inside a @contextmanager generator, so cleanup is skipped on error.
• Yielding more than once from a @contextmanager generator.
• Reusing a @contextmanager object for two with blocks.
• Returning a generator expression that reads from a file opened in the same with block.
• Forgetting that __enter__ must return the value as should bind; without a return it binds None.

Best practices

• Use with for every resource that must be released.
• Prefer @contextmanager over writing the two methods by hand.
• Always wrap the yield in try ... finally.
• Return False from __exit__ unless suppression is the deliberate purpose.
• Use ExitStack when the number of resources is dynamic.
• Name managers as verbs or phrases that read well after with.

Practice

1. Write a context manager that times a block and reports even when it raises.
2. Write one that temporarily changes a dictionary and restores it afterwards.
3. Write a class based manager that suppresses one named exception type.
4. Use ExitStack to open a list of files whose length is decided at runtime.
5. Implement an atomic write manager and prove the original file survives a mid-write failure.

Conclusion

A context manager ties cleanup to a block rather than to your discipline. Write one with @contextmanager, put the yield inside try ... finally, and return False from __exit__ unless swallowing the exception is exactly what you mean.