Decorators and Context Managers

In the Python for Beginners series you have already written functions and opened files with with. Chapter 19 of this 24-chapter course introduces the two tools that scale both habits: decorators bolt extra behaviour onto a function without editing its body, and context managers guarantee that setup and cleanup always run as a pair.

Functions are objects — decorators use that

A function is a first-class value: assign it, pass it, return it. A decorator is any callable that takes a function and hands back a replacement. The key is that loud(cheer) is just a normal call whose argument happens to be a function — and @loud is only sugar for writing that assignment. The inner wrapper is what actually runs: it calls the original, adds something, and returns the result. functools.wraps copies the original metadata onto the wrapper, so tools that read __name__ still see cheer:

import functools

def loud(func):
    @functools.wraps(func)
    def wrapper(*args, **kwargs):
        return func(*args, **kwargs).upper()
    return wrapper

@loud
def cheer(name):
    return f"Go {name}!"

print(cheer("Ada"))
print(cheer.__name__)

Expected output:

GO ADA!
cheer

Remove functools.wraps and the second line prints wrapper: debuggers, help() and test reports all start lying about which function they mean. An alias such as alias = greet and a decorator are the same idea — only the decorator gets to wrap the function on the way past.

Stacking decorators and giving them arguments

Decorators stack like layers: hello below is wrapped by bang first, then by tag("b"). A decorator that returns a decorator also gains arguments, which is why tag("b") needs parentheses — it is a function call whose result is the real decorator:

def tag(name):
    def decorator(func):
        def wrapper(*args, **kwargs):
            return f"<{name}>{func(*args, **kwargs)}</{name}>"
        return wrapper
    return decorator

def bang(func):
    def wrapper(*args, **kwargs):
        return func(*args, **kwargs) + "!"
    return wrapper

@tag("b")
@bang
def hello(name):
    return f"Hi {name}"

print(hello("Ada"))

Expected output: <b>Hi Ada!</b>

Decorators do not have to be functions. A class works too: __init__ receives the decorated function and a __call__ method does the work, which is handy when the decorator must remember state between calls — a call counter is the classic example. Use functools.update_wrapper(self, func) in __init__ to keep the original name, and @ treats the class exactly like a function.

Context managers: cleanup you can trust

A context manager is an object with two special methods. __enter__ runs when the with block starts and its return value is bound to as; __exit__ runs when the block ends — even when an exception tears through it. That is why with open(...) closes files, why with lock: releases a mutex (the file version appeared in chapter 15) and why timing wrappers never leak:

class LockedFile:
    def __init__(self, path):
        self.path = path

    def __enter__(self):
        print(f"opening {self.path}")
        self.file = open(self.path, "w", encoding="utf-8")
        return self.file

    def __exit__(self, exc_type, exc, tb):
        self.file.close()
        print(f"closed {self.path}")

with LockedFile("note.txt") as handle:
    handle.write("written safely")

Expected output:

opening note.txt
closed note.txt

__exit__ receives the exception type, value and traceback — or three Nones when the block succeeded. Return nothing and exceptions keep travelling upward; return True and they are swallowed, which is almost never what you want.

Writing a whole class for a small scope is verbose, so contextlib.contextmanager turns a generator into a context manager instead: lines before yield are the setup, the yielded value goes to as, and lines after yield are the cleanup. Wrap the yield in try/finally so cleanup still runs when the body raises:

import time
from contextlib import contextmanager

@contextmanager
def slow_step(label):
    print(f"start {label}")
    start = time.perf_counter()
    try:
        yield
    finally:
        print(f"end {label} (slept 0.1s: {time.perf_counter() - start >= 0.1})")

with slow_step("report"):
    time.sleep(0.1)

Expected output:

start report
end report (slept 0.1s: True)

Worked example: retrying a flaky call

Now combine everything: a decorator that takes an argument, preserves the wrapped name, catches the ConnectionError you met in chapter 12, and fails loudly when the attempts run out:

import functools

def retry(times):
    def decorator(func):
        @functools.wraps(func)
        def wrapper(*args, **kwargs):
            for attempt in range(1, times + 1):
                try:
                    return func(*args, **kwargs)
                except ConnectionError:
                    print(f"attempt {attempt} failed, retrying")
            raise RuntimeError(f"{func.__name__} gave up")
        return wrapper
    return decorator

attempts = {"n": 0}

@retry(times=3)
def fetch():
    attempts["n"] += 1
    if attempts["n"] < 3:
        raise ConnectionError("timeout")
    return "payload"

print("result:", fetch())

Expected output:

attempt 1 failed, retrying
attempt 2 failed, retrying
result: payload

The same shape fits timing, caching, logging or permission checks. Reach for with when something is acquired and released, and for a decorator when a call needs extra behaviour — each still does exactly one job, even when they surround each other.

Common mistakes

• Forgetting functools.wraps — every wrapper then reports __name__ == "wrapper", so debuggers and test output point at the wrong function. Decorate each inner wrapper.
• A decorator that forgets to return — miss return wrapper (or return decorator) and the decorated name silently becomes None; the first call raises TypeError: 'NoneType' object is not callable.
• Skipping try/finally after yield — in a @contextmanager generator the cleanup after yield never runs if the body raises, so your “close” silently stops happening under error conditions.
• Returning True from __exit__ — it suppresses exceptions and hides real failures. Only swallow an exception you have genuinely handled.
• Using the resource after the block — the file is closed and the lock released the moment the with block ends; finish the work inside it or copy the data out first.

Practice checklist

1. Write your own loud decorator, delete functools.wraps, and watch cheer.__name__ change.
2. Stack two decorators and predict the exact output of hello("Ada") before running it.
3. Apply @retry(times=3) to a function that always raises and confirm the RuntimeError appears after exactly three attempts.
4. Make the body of a with block raise an exception and prove your context manager still cleans up.
5. Convert a class-based context manager into an @contextmanager generator (or the reverse) and compare the two versions.

What's next?

Decorators wrap behaviour and context managers wrap scope; both work best when the contracts are written down. Next you write them literally — Type Hints in Python shows how to annotate functions, containers and optional values so a checker can read your intentions.