Home/Learn/Python A–Z/Context Managers — with, __enter__/__exit__ & contextlib

Context Managers — with, __enter__/__exit__ & contextlib

Advanced
Iterators & Decorators

The with statement guarantees setup/teardown around a block — files, locks, DB transactions — via __enter__/__exit__ or the @contextmanager generator shortcut.

Overview

Anything shaped "acquire, use, ALWAYS release" belongs in a context manager: files, database connections/transactions, locks, timers, temporary state changes. with calls __enter__ before the block and __exit__ after — even when the block raises. Writing your own is either a small class or, more elegantly, a generator with @contextmanager: code before yield is setup, after is teardown, and try/finally makes it exception-safe. FastAPI dependencies and pytest fixtures reuse exactly this idea.

How with Works & a Custom Class

__enter__ returns the value bound by as; __exit__ receives exception info (type, value, traceback) — return True to swallow the exception, anything falsy to propagate.

Transactions — the canonical enter/exit example
class Transaction:
    def __init__(self, db):
        self.db = db
    def __enter__(self):
        print("BEGIN")
        return self.db                     # bound to 'as' target
    def __exit__(self, exc_type, exc, tb):
        if exc_type is None:
            print("COMMIT")
        else:
            print(f"ROLLBACK ({exc})")
        return False                       # False = re-raise if error

class FakeDB:
    def save(self, x): print("saved", x)

try:
    with Transaction(FakeDB()) as db:
        db.save("order-1")
        raise ValueError("payment failed")
except ValueError:
    pass
# BEGIN / saved order-1 / ROLLBACK (payment failed)

with Transaction(FakeDB()) as db:
    db.save("order-2")
# BEGIN / saved order-2 / COMMIT

@contextmanager — the Generator Shortcut

One generator, one yield: before = __enter__, after = __exit__. Wrap the yield in try/finally so teardown survives exceptions. Much less boilerplate than the class.

yield splits setup from guaranteed teardown
from contextlib import contextmanager
import time, os

@contextmanager
def timer(label):
    start = time.perf_counter()
    try:
        yield                              # block runs here
    finally:
        ms = (time.perf_counter() - start) * 1000
        print(f"{label}: {ms:.1f} ms")

with timer("sorting"):
    sorted(range(1_000_000, 0, -1))
# sorting: 61.2 ms

@contextmanager
def env_var(key, value):                   # temporary state change
    old = os.environ.get(key)
    os.environ[key] = value
    try:
        yield
    finally:                               # ALWAYS restore
        if old is None:
            os.environ.pop(key, None)
        else:
            os.environ[key] = old

with env_var("MODE", "test"):
    print(os.environ["MODE"])              # test
print(os.environ.get("MODE"))              # back to original

Key Points to Remember

  • 1with = __enter__ before the block, __exit__ after — even on exceptions
  • 2__exit__ returning True swallows the exception; False propagates it
  • 3@contextmanager: setup before yield, teardown after, wrapped in try/finally
  • 4Use for anything acquire/release: files, locks, transactions, timers, temp state

Interview Questions

Sign in to ask Aria
1

What methods must a context manager implement, and what does each receive/return?

MediumWalmart
2

Write a context manager that times a code block, using @contextmanager.

MediumRazorpay
3

How can __exit__ suppress an exception? When is that appropriate?

HardAtlassian

Ask Aria about Context Managers — with, __enter__/__exit__ & contextlib

Your personal AI tutor — ask anything about this concept

Revision Status

Personal Notes

Sign in to save personal notes for this topic.

Discussion

Sign in to join the discussion.

Loading discussion…