Context Managers

The with Statement

A context manager guarantees cleanup happens — even on exceptions. The canonical example is open().

# Without with — fragile
f = open("data.txt")
try:
    data = f.read()
finally:
    f.close()

# With — concise and exception-safe
with open("data.txt") as f:
    data = f.read()
# f is guaranteed closed here, even if read() raises

Multiple Contexts

with open("in.txt") as src, open("out.txt", "w") as dst:
    dst.write(src.read())

Writing Your Own Context Manager

Class-based: __enter__ and __exit__

import time

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

    def __enter__(self):
        self.start = time.perf_counter()
        return self

    def __exit__(self, exc_type, exc, tb):
        elapsed = time.perf_counter() - self.start
        print(f"{self.label}: {elapsed*1000:.2f}ms")
        # Return True to swallow the exception (don't, normally)
        return False

with Timer("db query"):
    run_query()

Function-based: @contextmanager

from contextlib import contextmanager

@contextmanager
def timer(label):
    start = time.perf_counter()
    try:
        yield                       # body of `with` runs here
    finally:
        elapsed = time.perf_counter() - start
        print(f"{label}: {elapsed*1000:.2f}ms")

with timer("db query"):
    run_query()

contextlib Helpers

contextlib.closing

from contextlib import closing
from urllib.request import urlopen

with closing(urlopen(url)) as resp:
    data = resp.read()
# closing() makes any object with .close() into a context manager

contextlib.suppress

from contextlib import suppress

with suppress(FileNotFoundError):
    os.remove("tmp.txt")
# Equivalent to try/except FileNotFoundError: pass

contextlib.ExitStack

Dynamic numbers of context managers — open N files known at runtime.

from contextlib import ExitStack

with ExitStack() as stack:
    files = [stack.enter_context(open(p)) for p in paths]
    for f in files:
        process(f.read())
# All N files closed in reverse order

chdir (3.11+)

from contextlib import chdir

with chdir("/tmp"):
    do_work()
# cwd is restored on exit

Async Context Managers

from contextlib import asynccontextmanager
import aiohttp

@asynccontextmanager
async def api_session(token):
    session = aiohttp.ClientSession(headers={"Authorization": f"Bearer {token}"})
    try:
        yield session
    finally:
        await session.close()

async def main():
    async with api_session(token) as s:
        async with s.get("/users") as r:
            print(await r.json())

Real-World Patterns

Database transactions

@contextmanager
def transaction(conn):
    try:
        yield conn
        conn.commit()
    except:
        conn.rollback()
        raise

with transaction(conn) as tx:
    tx.execute("INSERT ...")
    tx.execute("UPDATE ...")
# commit on success, rollback on exception

Temporary monkey-patching for tests

@contextmanager
def override_setting(name, value):
    original = settings[name]
    settings[name] = value
    try:
        yield
    finally:
        settings[name] = original

with override_setting("DEBUG", True):
    run_test()