Context Managers
Some resources need cleanup.
- A file handle has to be closed.
- A database connection has to be released back to the pool.
- A lock has to be unlocked, even if the code holding it raises halfway through.
The with statement is Python's tool for tying setup and cleanup together so the cleanup always runs.
The with statement
Without with, opening and closing a file looks like this:
f = open("tasks.txt")
try:
contents = f.read()
finally:
f.close()
The try/finally is there so that close() runs even if read() raises.
Warning
Forget the finally and a raised exception leaves the file open until the garbage collector gets to it.
with collapses the same pattern into one block:
with open("tasks.txt") as f:
contents = f.read()
When the block exits, whether normally or because of an exception, Python calls the cleanup code on f, without needing to manually specify it.
The file is closed before the next line runs.
The __enter__ and __exit__ protocol
Anything with __enter__ and __exit__ methods can be used in a with block. __enter__ runs at the start, and whatever it returns gets bound to the name after as. __exit__ runs at the end and receives details about any exception that was raised inside the block.
class Section:
def __init__(self, label: str) -> None:
self.label = label
def __enter__(self) -> "Section":
print(f"--- {self.label} ---")
return self
def __exit__(self, exc_type, exc_value, traceback) -> None:
print(f"--- end {self.label} ---")
with Section("normalize tasks"):
print("doing work")
Writing a class for every context manager is more ceremony than most cases need. The contextlib module has a shortcut.
contextlib.contextmanager
@contextmanager lets you write a context manager as a generator function. Everything before yield is the setup, the yielded value is what gets bound to the as name, and everything after yield is the cleanup.
import time
from contextlib import contextmanager
@contextmanager
def timer(label: str):
start = time.perf_counter()
try:
yield
finally:
elapsed_ms = (time.perf_counter() - start) * 1000
print(f"{label}: {elapsed_ms:.2f}ms")
with timer("normalize tasks"):
normalized = [t["title"].strip().title() for t in raw_tasks]
The try/finally around yield is important.
Without it, a raised exception inside the with block would skip the cleanup. With it, the elapsed time prints whether the block succeeded or blew up.
Info
When we work with database sessions later, we'll use the same pattern: a with block opens a session, runs a few queries, and the cleanup commits or rolls back the transaction before closing the session. That keeps every database interaction scoped to a single block, instead of trusting every code path to remember to clean up.