Python Decorators vs Context Managers: When to Use Each

Decorators and context managers solve the same kind of problem: running some code before and after other code, such as starting a timer and reporting it, opening a resource and closing it, or beginning a transaction and committing or rolling it back. The difference is scope. A decorator wraps a whole function, every time it is called. A context manager wraps exactly the block of code inside a with statement. Once you think in those terms, most choices make themselves.

A context manager wraps a block

A with statement calls the object's __enter__ method before the block and its __exit__ method after it, however the block ends: normally, through return, or by raising an exception (PEP 343). The easiest way to write one is a generator function with @contextmanager: the code before yield is the setup, the code after it is the clean-up (Python docs: contextlib).

import time
from contextlib import contextmanager

@contextmanager
def timer(label):
    start = time.perf_counter()
    try:
        yield
    finally:
        print(f"{label}: {time.perf_counter() - start:.1f}s")

def nightly_job():
    rows = list(range(1000))            # fast preparation, not timed
    with timer("slow section"):
        time.sleep(0.1)                  # only this part is measured
    return len(rows)

print(nightly_job())

slow section: 0.1s
1000

The try/finally around yield is what guarantees the report even if the block raises. The key point is precision: the timer measures one section of the function, which a decorator cannot do.

Handling errors inside the block: a transaction

Context managers can react to exceptions. __exit__ receives the exception, if there was one, so it can undo work. Database transactions are the classic case: commit if the block succeeded, roll back if it failed. Python's own sqlite3 connection behaves this way when used in a with statement (Python docs: sqlite3):

import sqlite3

db = sqlite3.connect(":memory:")
db.execute("CREATE TABLE accounts (name TEXT PRIMARY KEY, balance INTEGER CHECK (balance >= 0))")
db.executemany("INSERT INTO accounts VALUES (?, ?)", [("asha", 500), ("ravi", 100)])
db.commit()

def transfer(amount):
    try:
        with db:                                   # commit on success, roll back on error
            db.execute("UPDATE accounts SET balance = balance + ? WHERE name = 'ravi'", (amount,))
            db.execute("UPDATE accounts SET balance = balance - ? WHERE name = 'asha'", (amount,))
    except sqlite3.IntegrityError as exc:
        print("rolled back:", exc)
    print(dict(db.execute("SELECT name, balance FROM accounts")))

transfer(200)
transfer(1000)

{'asha': 300, 'ravi': 300}
rolled back: CHECK constraint failed: balance >= 0
{'asha': 300, 'ravi': 300}

In the failed transfer, Ravi's balance was already credited when Asha's debit broke the constraint. The rollback undid the credit, so no money appeared from nowhere. The transaction's scope is a few lines, not the whole function, which is exactly what a with block expresses.

A decorator wraps every call

When the behaviour applies to the entire function on every call, such as timing each request, retrying, caching or checking permissions, a decorator is the natural fit, and the rule is stated once at the top instead of inside the body. Python decorators with real examples builds several.

One object, both ways: ContextDecorator

Often you want both: time a whole function sometimes, and one block at other times. Context managers created with @contextmanager are also usable as decorators, because they are built on ContextDecorator (Python docs: contextlib). The timer above already works both ways:

@timer("whole export")
def export_all():
    time.sleep(0.1)
    return "done"

print(export_all())
with timer("ad hoc check"):
    time.sleep(0.1)

whole export: 0.1s
done
ad hoc check: 0.1s

As a decorator, a fresh context is entered for each call. This is the cleanest answer to "decorator or context manager?" when a behaviour genuinely has both uses: write the context manager, and get the decorator for free.

When to write a class instead

@contextmanager covers most cases. Write a class with __enter__ and __exit__ when the context manager needs state that outlives the block or methods of its own, for example a connection pool object that is also used outside with. Two details matter. Whatever __enter__ returns is what as binds, so return self if the caller should use the object inside the block. And __exit__(exc_type, exc, tb) should normally return None or False, letting any exception continue after the clean-up; returning True swallows it, which is almost never what the caller expects.

Asynchronous code has matching tools: async with calls __aenter__ and __aexit__, and contextlib.asynccontextmanager builds one from an async def generator. Async database sessions and HTTP clients are usually used this way.

Two contextlib helpers worth knowing

suppress states that a specific exception is expected and harmless, more plainly than an empty except block:

import os
from contextlib import ExitStack, suppress

with suppress(FileNotFoundError):
    os.remove("old-cache.tmp")          # fine if it was never there
print("continued")

continued

ExitStack manages a number of context managers that is only known at run time, such as opening every file in a list, and closes all of them, in reverse order, even if opening one fails halfway:

names = ["a.txt", "b.txt", "c.txt"]
for n in names:
    with open(n, "w", encoding="utf-8") as f:
        f.write(n.upper())

with ExitStack() as stack:
    files = [stack.enter_context(open(n, encoding="utf-8")) for n in names]
    print([f.read() for f in files])
print(all(f.closed for f in files))

['A.TXT', 'B.TXT', 'C.TXT']
True

Choosing

Two short rules cover most cases. Resources and sections use a with statement. Files, locks, database transactions and temporary settings belong in context managers, because their lifetime is a block of code. Policies on whole functions use a decorator. Caching, retries, timing, registration and access control belong in decorators, because they apply to every call. When something is both, write a context manager with @contextmanager and use it either way.

Context managers built with @contextmanager are generators underneath; Python generators and iterators made simple explains yield. Clean-up on failure is part of good error handling, covered in error handling in Python. Decorators, closures and context managers are part of the advanced concepts module of the Python Programming course.

Sources

• contextlib — Utilities for with-statement contexts — Python documentation
• PEP 343 – The "with" Statement — Python Enhancement Proposals
• sqlite3: how to use the connection context manager — Python documentation