Decorators: Wrapping Functions Without Rewriting Them

Add behavior on the outside, leave the inside alone

Advanced 14 min

In this lesson

In Unit 9 you added timing and retries by editing the functions themselves. Do that six times and you have six copies of the same code around six different bodies. A decorator lets you write that wrapping once and apply it with a single line.

Explain it like I’m 5

A decorator is gift wrap. The present inside is unchanged; the wrapping just adds something on the outside.

The problem first

Here is the code a decorator replaces. Two functions, each with the same four lines of timing bolted around real work. Imagine it across a whole module.

Example 
import time

def fetch_users():
    start = time.perf_counter()      # copy 1
    result = ["ada", "grace"]
    print("fetch_users took", time.perf_counter() - start)
    return result

def fetch_orders():
    start = time.perf_counter()      # copy 2, identical
    result = [1, 2, 3]
    print("fetch_orders took", time.perf_counter() - start)
    return result
The real work is one line each. The rest is duplication.

A function that takes a function

In Python, functions are ordinary values. You can pass one to another function, and you can return one. A decorator is just that: it takes a function, builds a new function that adds something around it, and returns the new one.

The @ line is shorthand. @timer above def fetch(): means exactly fetch = timer(fetch).

Example 
def shout(func):
    def wrapper():
        print("-- before --")
        result = func()          # call the original
        print("-- after --")
        return result
    return wrapper               # hand back the new function

@shout
def greet():
    print("Hello!")

greet()
Output 
-- before --
Hello!
-- after --
The three-line shape: define wrapper, call the original, return wrapper.

Letting arguments through

The wrapper above takes no arguments, so it only decorates functions that take none. To wrap any function, use *args and **kwargs from Unit 4 to accept whatever comes and pass it straight through.

Add functools.wraps too. Without it the decorated function reports the wrapper's name and loses its docstring, which breaks help text and confuses debugging.

Example 
import functools

def shout(func):
    @functools.wraps(func)                    # keep the original identity
    def wrapper(*args, **kwargs):             # accept anything
        print("calling", func.__name__)
        return func(*args, **kwargs)          # pass it all through
    return wrapper

@shout
def add(a, b):
    """Add two numbers."""
    return a + b

print(add(2, 3))
print(add.__name__, "|", add.__doc__)
Output 
calling add
5
add | Add two numbers.
*args/**kwargs make it universal; wraps keeps the identity.

Build a timer decorator. Record how long the wrapped function took into DURATIONS under its name, and return the original result unchanged. (We store the duration rather than printing it, because a printed time would differ on every run.)

import time

DURATIONS = {}

def timer(func):
    def wrapper(*args, **kwargs):
        pass  # TODO: time the call, save to DURATIONS[func.__name__], return the result
    return wrapper

@timer
def add(a, b):
    return a + b

print(add(2, 3))
print("timed:", sorted(DURATIONS))
print("recorded a number:", isinstance(DURATIONS["add"], float))

Decorators that take arguments

@retry(times=3) needs one more layer. The outer function takes the settings and returns the actual decorator, which then takes the function. Three levels looks alarming written down and is mechanical once you have seen it once.

This is the retry loop from Unit 9, written once instead of at every call site.

Example 
import functools

ATTEMPTS = []

def retry(times):                      # 1. takes the settings
    def decorator(func):               # 2. takes the function
        @functools.wraps(func)
        def wrapper(*args, **kwargs):  # 3. does the work
            for attempt in range(1, times + 1):
                try:
                    return func(*args, **kwargs)
                except ValueError:
                    ATTEMPTS.append(attempt)
                    if attempt == times:
                        raise          # out of tries
        return wrapper
    return decorator

calls = []

@retry(times=3)
def flaky():
    calls.append(1)
    if len(calls) < 3:
        raise ValueError("not yet")
    return "worked on attempt 3"

print(flaky())
print("failed attempts:", ATTEMPTS)
Output 
worked on attempt 3
failed attempts: [1, 2]
Three layers: settings, function, call.

When not to write one

A decorator adds a layer of indirection: the reader now has to look somewhere else to know what a call does. That is worth it when the behavior genuinely repeats across many functions, and not worth it for something used once.

A workable rule: write it plainly the first two times. Wrap it when the same wrapping appears a third time.

Common mistake: Forgetting to return wrapper

Why it happens:

The wrapper is defined, so the job feels finished.

How to fix it:

The decorator must return the new function. Without it the decorated name becomes None and calling it raises TypeError.

Common mistake: Leaving out functools.wraps

Why it happens:

Everything works, so nothing prompts you to add it.

How to fix it:

The function loses its name and docstring, which breaks help() and makes tracebacks confusing. Add @functools.wraps(func) to every wrapper.

Common mistake: Defining wrapper without *args and **kwargs

Why it happens:

The first function you decorate happens to take no arguments.

How to fix it:

Accept *args, **kwargs and forward them, so the decorator works on any function.

Common mistake: Writing a decorator for one function

Why it happens:

It is a satisfying tool and everything starts looking wrappable.

How to fix it:

Put the code in the function. A decorator used once is indirection with no payoff.

What does @my_decorator above a function definition actually do?

Why does wrapper need *args and **kwargs?

What does functools.wraps preserve?

Mini exercise (medium)

Write a count_calls decorator that records how many times each decorated function has been called, in a module-level CALLS dict keyed by function name. It must work on functions with any arguments and preserve the original name and docstring.

Practice here. Fill in the missing piece and click Run to try your answer in place.

import functools

CALLS = {}

def count_calls(func):
    # TODO: wrap func so each call is counted in CALLS by name,
    #       keeping the original name and docstring
    return func

@count_calls
def greet(name):
    """Say hello."""
    return f"Hello, {name}!"

@count_calls
def add(a, b):
    return a + b

greet("Sam"); greet("Ada"); add(1, 2)
print(sorted(CALLS.items()))
print(greet.__name__, "|", greet.__doc__)

What to learn next

You wrote decorators from the ground up: a function taking a function and returning a new one, @ as shorthand for reassignment, *args/**kwargs so one decorator fits any signature, functools.wraps to keep the original identity, and the three-layer shape for decorators that take arguments. The retry loop from Unit 9 is now written once.

One repetition left. Context Managers and the with Statement handles the cleanup you keep remembering to do — and the times you forget.