intermediate30 minby DevFox

Context Managers

Guaranteeing cleanup even when something raises: the with statement, contextlib, writing your own, and the resource leaks that appear the moment you do not.

Python

Your service has been running for three days when it starts refusing connections. Nothing is wrong with the database, the network is fine, and the error is OSError: [Errno 24] Too many open files. Somewhere in a code path that runs occasionally, something is opened and never closed, and it has taken three days to run out.

Cleanup that depends on remembering is cleanup that eventually does not happen. By the end of this lesson you will write your own context managers, know the two ways to do it and when each fits, and understand the one design decision that turns a context manager into a bug.

The guarantee

You have used with on files:

python

with open(path, encoding="utf-8") as file:
    process(file)

The guarantee is precise: the cleanup runs when the block ends, however it ends. Normal completion, an exception, a return from inside, a break — all of them.

That "however" is the value. The version you can write by hand is longer and easier to get subtly wrong:

python

file = open(path, encoding="utf-8")
try:
    process(file)
finally:
    file.close()

Three lines of scaffolding around one line of work, repeated at every use, and the day someone adds an early return above the finally, it still works — but the day someone writes the open outside the try, it does not.

A context manager packages that pattern once so callers get it for free.

Writing one with a generator

The short way uses contextlib and the generator syntax from the previous lesson:

python

from contextlib import contextmanager

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

python

with timer("captioning"):
    caption_everything(photos)
# captioning took 4.71s

Read it as three parts. Everything before the yield is setup, run on entry. The yield is where the body of the with block happens. Everything after is cleanup.

The try/finally is not optional. Without it, an exception in the body propagates from the yield and the cleanup never runs — which is exactly the failure the context manager exists to prevent.

To hand something to the block, yield it:

python

@contextmanager
def open_database(path):
    connection = sqlite3.connect(path)
    try:
        yield connection
    finally:
        connection.close()

python

with open_database("photos.db") as connection:
    connection.execute("SELECT * FROM photos")

Whatever you yield is what as binds.

Writing one as a class

The other way is two methods:

python

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

    def __enter__(self):
        self.start = time.perf_counter()
        return self                      # what `as` binds

    def __exit__(self, exc_type, exc_value, traceback):
        self.elapsed = time.perf_counter() - self.start
        print(f"{self.label} took {self.elapsed:.2f}s")
        return False                     # do not suppress

python

with Timer("captioning") as t:
    caption_everything(photos)

print(t.elapsed)          # still available afterwards

__enter__ runs on entry and returns what as binds. __exit__ runs on the way out and receives details of the exception, if there was one — all three arguments are None on a clean exit.

Choose the class form when the manager needs to hold state the caller reads afterwards, when it needs to be reusable or reentrant, or when it is a natural method on an object you already have. Choose the generator form otherwise, which is most of the time.

The return value of exit

This is the decision that turns a context manager into a bug.

return True

Suppresses the exception.

The block's failure is swallowed and the program continues after the with as though nothing happened.

Every caller of every with on this class now silently ignores errors, and nothing at the call site says so.

return False, or nothing

Cleans up and lets it through.

A missing return gives None, which is falsy — so the correct behaviour is also the one you get by not thinking about it.

One line, and it changes whether every failure inside every with block reaches its caller

Bad — a cleanup that silently eats every failure.

python

class Database:
    def __exit__(self, exc_type, exc_value, traceback):
        self.connection.close()
        return True             # "handled"

Good — cleanup that lets the failure through.

python

class Database:
    def __exit__(self, exc_type, exc_value, traceback):
        self.connection.close()
        return False            # or just return nothing

The first version closes the connection correctly and then tells Python the exception was dealt with. So a failed write inside the block produces no exception, no log line and no traceback — the caller believes the save succeeded, and the data is not there. One return True converts a loud failure into permanent silent data loss, and because everything else about the class is correct, it reviews well.

Return False, or return nothing at all, unless suppressing is the manager's entire declared purpose — which is rare, and contextlib.suppress already covers the legitimate case:

python

from contextlib import suppress

with suppress(FileNotFoundError):
    Path("cache.json").unlink()

That reads as what it does, at the call site, where a reader can see it.

Reacting to the exception

__exit__ can inspect what went wrong without suppressing it, which is how transactions work:

python

@contextmanager
def transaction(connection):
    try:
        yield connection
        connection.commit()          # only if the block succeeded
    except Exception:
        connection.rollback()
        raise                        # and let it continue upward

python

with transaction(connection) as tx:
    tx.execute("INSERT INTO photos ...")
    tx.execute("UPDATE counts ...")   # if this fails, both undo

The bare raise is doing the important work. The manager reacts — rolls back — and then gets out of the way. Removing that one line would make every failed transaction look successful.

Note the shape: commit after the yield in the success path, rollback in the handler. That is a different structure from try/finally, because the cleanup differs depending on what happened.

Several at once

Multiple managers in one statement, evaluated left to right and unwound right to left:

python

with open(source) as reader, open(target, "w") as writer:
    writer.write(reader.read())

Parentheses let you wrap it, which matters at 78 columns:

python

with (
    open(source, encoding="utf-8") as reader,
    open(target, "w", encoding="utf-8") as writer,
):
    ...

When the number is not known until runtime, ExitStack handles it:

python

from contextlib import ExitStack

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

What else this is for

Files and connections are the obvious cases. The pattern fits anything with a matching pair of actions:

text

open / close            files, sockets, connections
acquire / release       locks
begin / commit-rollback transactions
change / restore        working directory, settings, env vars
start / stop            timers, spans, profilers
enter / exit            a logging context, a test fixture

The temporarily-change-something case is worth showing, because the finally is what makes it correct:

python

@contextmanager
def temporary_setting(config, key, value):
    original = config[key]
    config[key] = value
    try:
        yield
    finally:
        config[key] = original       # restored even on failure

Without the try/finally, one exception leaves the setting changed for the rest of the program's life — a bug that shows up somewhere else entirely.

A cheat sheet to keep

text

THE GUARANTEE
  with thing() as x:
      ...
  cleanup runs however the block ends:
  normally, on an exception, on return, on break

GENERATOR FORM - the usual choice
  @contextmanager
  def managed():
      setup()
      try:
          yield value          <- what `as` binds
      finally:
          cleanup()
  the try/finally is NOT optional

CLASS FORM - when it holds state, or is reusable
  def __enter__(self):  return what `as` binds
  def __exit__(self, exc_type, exc_value, traceback):
      cleanup
      return False       <- or nothing

THE TRAP
  a TRUTHY return from __exit__ SUPPRESSES the exception
  the caller then believes it succeeded
  return False, or nothing, unless suppressing IS the purpose
  for the legitimate case: with suppress(FileNotFoundError):

REACTING WITHOUT SUPPRESSING
  except Exception:
      rollback()
      raise            <- the line that keeps it honest

SEVERAL
  with open(a) as f, open(b) as g:
  with ( ... , ... ):              wrap at 78 columns
  ExitStack()                      when the count is dynamic

WHAT IT IS FOR
  open/close, acquire/release, begin/commit, change/restore,
  start/stop - any pair where the second must always happen

  write try/finally twice -> write a context manager

Where to go next

You can now guarantee cleanup rather than remember it, in both forms, and you know the return True that quietly converts failures into silence. The instinct to keep: any pair of actions where the second must always happen is a context manager waiting to be written, and the duplication of try/finally is the signal.

Next is Decorators, which is the same idea applied to functions rather than blocks — wrapping behaviour around something without editing it. The generator-based context manager you just wrote is itself made by a decorator, so you have already used one to build one.

Before you move on, write a context manager for something in your own code that currently uses try/finally, and then write one that deliberately returns True from __exit__. Raise an exception inside its block and watch the program carry on regardless. That silence is the thing to remember.