Lesson 11. Context Managers

Mission link: Every resource that has to be released, every setting that has to be restored, and every temporary state that has to be undone is the same problem, and Python has one construct for all of them. Recognising the pattern is what stops try/finally from spreading.
Primary source: The Python Language Reference, The with statement
Prerequisites: Lesson 9, Lesson 10

Warm-up

1. ▢ Lesson 10: what does finally guarantee that an except clause does not?

2. ▢ Lesson 9 put a with open(...) inside a generator. When does the file close if the caller abandons the generator halfway?

Know this

with is try/finally with a name. This:

with open(path) as f:
    data = f.read()

is this:

manager = open(path)
f = manager.__enter__()
try:
    data = f.read()
finally:
    manager.__exit__(*sys.exc_info())

Two methods make an object usable this way.

__enter__(self) runs on entry, and its return value is what as binds. That is worth stating separately, because it is not necessarily the manager. A file returns itself, so f is the file. threading.Lock.__enter__ returns the result of acquire, so with lock as x binds True, which is why nobody writes the as there.

__exit__(self, exc_type, exc, tb) runs on exit. The three arguments are None, None, None on a clean exit and describe the exception otherwise. Returning a truthy value suppresses that exception; returning None lets it propagate, which is what almost every __exit__ should do.

Writing one with a generator

The everyday way is a generator plus a decorator. Everything before the yield is entry, everything after is exit, and the yields value is what as binds:

from contextlib import contextmanager

@contextmanager
def timed(label):
    start = time.perf_counter()
    try:
        yield                          # the body of the with block runs here
    finally:
        log.info("%s took %.3fs", label, time.perf_counter() - start)

with timed("import"):
    run_import()

The try/finally inside is not optional. Without it, an exception in the caller's block propagates out of the yield and the cleanup never runs. That is the one thing to get right in a generator-based manager, and forgetting it produces a manager that works in every test and leaks in production.

To bind something, yield it:

@contextmanager
def temporary_table(conn, name):
    conn.execute(f"CREATE TEMP TABLE {name} (id bigint)")
    try:
        yield name
    finally:
        conn.execute(f"DROP TABLE {name}")

A generator-based manager is single-use. Each with needs a fresh call, because the generator it built is an iterator and lesson 9 applies. When a manager must be entered twice, or nested inside itself, write a class.

Writing one as a class

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

    def __enter__(self):
        self.conn.execute("BEGIN")
        return self.conn

    def __exit__(self, exc_type, exc, tb):
        if exc_type is None:
            self.conn.execute("COMMIT")
        else:
            self.conn.execute("ROLLBACK")
        return None                    # never swallow: the caller must know

The exc_type is None test is the reason to reach for a class: entry and exit differ by outcome, not just by cleanup. Note the explicit return None, and note how easy the bug is: return True here would roll back and then tell the caller everything succeeded.

The first two steps are identical in both columns. The rollback is not what changes; what changes is whether anyone downstream is allowed to know it happened.

Several at once

with open(src) as fin, open(dst, "w") as fout:
    fout.write(fin.read())

Left to right on entry, right to left on exit, and the second open is protected by the first, so a failure there still closes fin. Parenthesised across lines is legal from Python 3.10:

with (
    open(src) as fin,
    open(dst, "w") as fout,
):
    ...

When the number is not known until runtime, ExitStack collects them:

from contextlib import ExitStack

with ExitStack() as stack:
    files = [stack.enter_context(open(p)) for p in paths]
    merge(files)                       # every file closed on the way out

ExitStack also solves the conditional case, where a resource is only sometimes needed, without an if around the whole block.

What contextlib already contains

nullcontext is the one that removes real duplication:

with (open(path) if path else nullcontext(sys.stdin)) as f:
    process(f)

The pattern, stated once

Anything with a paired operation belongs in a context manager: open and close, acquire and release, begin and commit, set and restore, mkdir and remove, patch and unpatch, connect and disconnect. If the pair is spelled out at two different indentation levels in a function, a manager will delete more code than it adds.

Async has its own pair, __aenter__ and __aexit__, used through async with. Stage 6 covers it. The rules above transfer unchanged.

Practice

1. ▢ What is wrong with this manager, and when does it show up?

   @contextmanager
   def locked(lock):
       lock.acquire()
       yield
       lock.release()
   

@contextmanager
def locked(lock):
    lock.acquire()
    try:
        yield
    finally:
        lock.release()

2. ▢ Predict the output.

   class Quiet:
       def __enter__(self):
           return self
       def __exit__(self, exc_type, exc, tb):
           return True
   
   with Quiet():
       print("before")
       raise ValueError("boom")
       print("after")
   print("done")
   

before
done

3. ▢ What does x hold?

   with open("notes.txt") as f, threading.Lock() as x:
       ...
   

4. ▢ Rewrite without try/finally.

   original = os.environ.get("TZ")
   os.environ["TZ"] = "UTC"
   try:
       run_report()
   finally:
       if original is None:
           del os.environ["TZ"]
       else:
           os.environ["TZ"] = original
   

@contextmanager
def env(name, value):
    original = os.environ.get(name)
    os.environ[name] = value
    try:
        yield
    finally:
        if original is None:
            del os.environ[name]
        else:
            os.environ[name] = original

with env("TZ", "UTC"):
    run_report()

5. ▢ This opens a variable number of files. Why does the naive version leak, and what fixes it?

   files = [open(p) for p in paths]
   try:
       merge(files)
   finally:
       for f in files:
           f.close()
   

with ExitStack() as stack:
    files = [stack.enter_context(open(p)) for p in paths]
    merge(files)

6. ▢ Name two things in code you have seen that are paired operations written without a context manager.

Real-world reps

• [ ] Find a try/finally in code you own whose finally restores something rather than closing something. Turn it into a context manager and convert two call sites.
• [ ] Write the timed manager from this lesson from memory, including the try/finally, and use it on something slow you already have.
• [ ] Find a place that opens a resource and closes it in a different function, so the pairing is invisible. Decide whether a manager can bring them together, and write down why if it cannot.
• [ ] Tomorrow: read the contextlib page once. It is a list of problems you have already solved by hand.

Going further

• The with statement: the exact desugaring, including the multiple-manager form
• With Statement Context Managers: __enter__ and __exit__ as the data model states them
• contextlib: the decorator, ExitStack, and the ready-made managers
• PEP 343, The "with" Statement: the rationale, and why suppression is opt-in
• Resources

Not landing? Reread the primary source at the top, since this lesson compresses it and compression is where understanding leaks. Check the glossary for any term that felt slippery.

If the lesson itself is unclear rather than the material, that is a defect: open an issue.