Context Managers (reliable setup/teardown)

When to use

  • You need “do X, then always clean up”: transactions, temp files/dirs, locks, creds.
  • The cleanup must run even if errors happen.
  • You’re repeating the same try/except/finally pattern.

Avoid when there’s no resource to release or one plain function call is enough.

Diagram (text)

with transaction(db) as cur:
    cur.execute(...)
# __enter__ → start work
# __exit__  → commit or rollback on error

Python example (≤40 lines, type-hinted)

Simple, productiony DB transaction with tests using a fake DB.

from __future__ import annotations
from contextlib import contextmanager
from typing import Protocol, Iterator

class Cursor(Protocol):
    def execute(self, sql: str) -> None: ...

class DB(Protocol):
    def cursor(self) -> Cursor: ...
    def commit(self) -> None: ...
    def rollback(self) -> None: ...

@contextmanager
def transaction(db: DB) -> Iterator[Cursor]:
    cur = db.cursor()
    try:
        yield cur                # caller runs work here
    except Exception:
        db.rollback()            # always clean up on error
        raise
    else:
        db.commit()              # only if no exception

# --- tiny fakes + tests ---
class FakeCur:
    def __init__(self, log): self.log = log
    def execute(self, sql: str) -> None: self.log.append(sql)

class FakeDB:
    def __init__(self): self.log=[]; self.commits=0; self.rollbacks=0
    def cursor(self) -> Cursor: return FakeCur(self.log)
    def commit(self) -> None: self.commits += 1
    def rollback(self) -> None: self.rollbacks += 1

def test_transaction_commit_and_rollback():
    db = FakeDB()
    with transaction(db) as cur: cur.execute("INSERT ok")
    assert db.commits==1 and db.rollbacks==0 and db.log==["INSERT ok"]

    db2 = FakeDB()
    import pytest
    with pytest.raises(ValueError):
        with transaction(db2) as cur: cur.execute("INS bad"); raise ValueError
    assert db2.commits==0 and db2.rollbacks==1

Trade-offs & pitfalls

  • Pros: Cleanup is guaranteed; removes boilerplate; easy to test; safer error paths.
  • Cons: Control flow is a bit hidden; long with scopes can hold resources too long.
  • Pitfalls:
    • Forgetting to re-raise exceptions in __exit__/except block.
    • Doing heavy work in __enter__ that can fail half-initialized.
    • Nested transactions without savepoints; design policy explicitly.
    • For async libs, you must use async with variants.

Pythonic alternatives

  • contextlib.contextmanager (used above) or class-based __enter__/__exit__.
  • contextlib.ExitStack to manage multiple resources together.
  • contextlib.asynccontextmanager for async DB/HTTP clients.
  • Many drivers already expose with db.transaction(): ...—use library-provided managers when available.

Mini exercise

Add a savepoint(db, name) context manager that logs "SAVEPOINT name" on enter, "RELEASE SAVEPOINT name" on success, and "ROLLBACK TO SAVEPOINT name" on error (then re-raise). Extend FakeDB to record these calls and test nested transactions.

Checks (quick checklist)

  • Cleanup runs on both success and error paths.
  • Exceptions are re-raised after rollback.
  • Scope of the with is minimal (don’t hold resources longer than needed).
  • Tests cover success and failure.
  • Use async with / ExitStack when appropriate.