Context Managers & with

Guarantee setup and cleanup with the with statement — via __enter__/__exit__ or @contextmanager.

The with statement guarantees that cleanup happens — even if an exception is raised. You already use it for files: with open(path) as f: closes the file no matter what. You can give your own objects the same superpower.

The two ways to build one

1. The class protocol. Implement __enter__ (runs on entry, its return value is bound to as) and __exit__ (always runs on the way out, even on error):

class Timer:
    def __enter__(self):
        self.start = time.perf_counter()
        return self
    def __exit__(self, exc_type, exc, tb):
        self.elapsed = time.perf_counter() - self.start
        # return False to let any exception propagate

2. The decorator way. @contextlib.contextmanager turns a generator into a context manager: everything before yield is setup, everything after is teardown, and the yielded value is what as receives.

The exit return value

If __exit__ returns a truthy value, it swallows the exception. Return False (or nothing) to let errors propagate — that is almost always what you want.

Example

import time
from contextlib import contextmanager

# The generator style — setup, yield, guaranteed teardown
@contextmanager
def timed(label):
    start = time.perf_counter()
    print(f'[{label}] start')
    try:
        yield                       # body of the with-block runs here
    finally:
        # runs even if the block raises
        dt = (time.perf_counter() - start) * 1000
        print(f'[{label}] done in {dt:.1f}ms')

# The class style — a reusable transaction-like manager
class Transaction:
    def __init__(self):
        self.log = []
    def __enter__(self):
        self.log.append('BEGIN')
        return self
    def __exit__(self, exc_type, exc, tb):
        self.log.append('ROLLBACK' if exc_type else 'COMMIT')
        return False                # never swallow the error

with timed('work'):
    sum(range(100_000))

tx = Transaction()
try:
    with tx:
        raise ValueError('boom')
except ValueError:
    pass
print('transaction log:', tx.log)   # cleanup ran despite the error

# Output (timing will vary):
# [work] start
# [work] done in 2.4ms
# transaction log: ['BEGIN', 'ROLLBACK']

When to use it

A transaction context manager commits on success and rolls back on exception without try/finally boilerplate.
A temporary file context manager creates a file, yields its path, and deletes it when the block exits.
A test isolation helper uses contextlib.ExitStack to open a variable number of context managers in one with.

More examples

contextmanager for transactions

Implements a commit/rollback transaction context with @contextmanager and a try/except/yield.

from contextlib import contextmanager

@contextmanager
def transaction(conn):
    try:
        yield conn
        conn['commit'] = True
        print('Committed')
    except Exception:
        conn['rollback'] = True
        print('Rolled back')
        raise

conn = {}
with transaction(conn):
    conn['sql'] = 'INSERT ...'
print(conn)   # {'commit': True, 'sql': 'INSERT ...'}

Suppress specific exceptions

Uses contextlib.suppress() to silently ignore a FileNotFoundError instead of try/except/pass.

from contextlib import suppress
import os

with suppress(FileNotFoundError):
    os.remove('nonexistent.tmp')
print('Continues without error')

contextlib.ExitStack for dynamic cleanup

Uses ExitStack to open a variable number of files and guarantee they all close when the block exits.

from contextlib import ExitStack

filenames = ['a.txt', 'b.txt']
with ExitStack() as stack:
    files = [stack.enter_context(open(f, 'w'))
             for f in filenames]
    files[0].write('hello')
    files[1].write('world')
print('All files closed')

Discussion

Be the first to comment on this lesson.