Decorators

What you'll learn

• Write a decorator: a function that takes a function and returns a wrapped version
• Understand that @name is only syntax for func = name(func)
• Keep the wrapped function's name and docstring with functools.wraps
• Write decorators that take arguments, and stack more than one

Concept

A decorator is a function that receives a function and returns a new function that usually calls the original with something added: timing, logging, caching, access checks, retries. You have used them: @property, @classmethod, @staticmethod, @abstractmethod. Now you will write them, and they turn out to be closures with a convention on top.

The @ line is pure syntax. Writing @timed above def slow(): means exactly slow = timed(slow) after the definition. The decorator runs once, at definition time, and whatever it returns is bound to the name. Understanding that one line of equivalence removes all the mystery: a decorator is a function call, and the result replaces the original.

The standard shape has three levels. The decorator takes func. Inside it defines wrapper(*args, **kwargs), which does something, calls func(*args, **kwargs), does something else, and returns the result. The decorator returns wrapper. The star parameters let the wrapper accept whatever the original accepts, and returning the result keeps the original's return value flowing through. Forget to return func's result and every decorated function silently returns None.

Without help, the wrapper has its own name and no docstring, so help(slow) shows wrapper and debugging output lies. functools.wraps(func) applied to the wrapper copies the original's metadata across; it belongs in every decorator you write.

When a decorator needs its own arguments, @retry(times=3), you add one more level: a function that takes the arguments and returns a decorator. @retry(times=3) calls retry(3) first, then applies the result exactly as before.

Decorators stack. @a above @b above def f means f = a(b(f)): the one nearest the function is applied first. Keep decorators small and single-purpose, so stacking reads as a list of properties: timed, cached, logged. Class methods can be decorated the same way; the wrapper receives self as its first positional argument.

Code

The simplest decorator, written out with and without the @ syntax, to make the equivalence concrete.

# equivalence.py
def shout(func):
    def wrapper(*args, **kwargs):
        result = func(*args, **kwargs)
        return result.upper()
    return wrapper


def greet(name):
    return f"hello, {name}"


greet = shout(greet)
print(greet("ada"))


@shout
def farewell(name):
    return f"bye, {name}"


print(farewell("linus"))

HELLO, ADA
BYE, LINUS

A useful decorator with functools.wraps: timing. Compare the two decorated functions' metadata to see what wraps preserves.

# timed.py
import functools
import time


def timed(func):
    @functools.wraps(func)
    def wrapper(*args, **kwargs):
        start = time.perf_counter()
        result = func(*args, **kwargs)
        ms = (time.perf_counter() - start) * 1000
        print(f"{func.__name__} took under {int(ms) // 50 * 50 + 50} ms")
        return result
    return wrapper


def naive(func):
    def wrapper(*args, **kwargs):
        return func(*args, **kwargs)
    return wrapper


@timed
def total(n):
    """Sum the first n integers."""
    return sum(range(n))


@naive
def other():
    """Does nothing."""


print(total(1_000_000))
print(total.__name__, "|", total.__doc__)
print(other.__name__, "|", other.__doc__)

total took under 50 ms
499999500000
total | Sum the first n integers.
wrapper | None

A decorator with arguments: retry a flaky function up to times attempts. The extra level receives the argument and returns the actual decorator.

# retry.py
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 as err:
                    print(f"  attempt {attempt} failed: {err}")
            raise RuntimeError(f"{func.__name__} failed after {times} attempts")
        return wrapper
    return decorator


calls = {"n": 0}


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


print(fetch())
print("calls made:", calls["n"])

  attempt 1 failed: timeout
  attempt 2 failed: timeout
data
calls made: 3

Stacking, and decorating a method. The innermost decorator applies first, and the wrapper on a method receives self in args.

# stack.py
import functools


def log_calls(func):
    @functools.wraps(func)
    def wrapper(*args, **kwargs):
        print(f"-> {func.__name__}{args[1:] if args and hasattr(args[0], '__dict__') else args}")
        return func(*args, **kwargs)
    return wrapper


def double(func):
    @functools.wraps(func)
    def wrapper(*args, **kwargs):
        return func(*args, **kwargs) * 2
    return wrapper


@log_calls
@double
def add(a, b):
    return a + b


print(add(2, 3))


class Cart:
    def __init__(self):
        self.items = []

    @log_calls
    def add(self, item):
        self.items.append(item)
        return len(self.items)


print(Cart().add("pen"))

-> add(2, 3)
10
-> add('pen',)
1

Try it yourself

Create decorate.py with three decorators: debug, which prints the function name and arguments before calling and the result after; count_calls, which keeps a per-function call count on the wrapper as an attribute; and require_positive, which raises ValueError if any argument is not positive. Apply all three to one function, call it several times, and print the call count. Then write @repeat(n) that runs the function n times and returns the last result.

Common mistakes

• Forgetting return inside the wrapper. The original's result is lost and every call returns None.
• Skipping functools.wraps. Tracebacks, help, and logging then show wrapper for every decorated function.
• Adding arguments without the extra level. @retry and @retry(3) are different shapes; the second needs a function that returns a decorator.

What's next

• @d above def f means f = d(f); a decorator is a closure that wraps a function and returns the wrapper.
• Use *args, **kwargs, return the original's result, and apply functools.wraps in every decorator.
• Decorators with arguments add one more level; stacked decorators apply from the bottom up.

You can now add behavior to any function without touching it. Next you will meet the functional tools that pair with decorators: lambda, map, filter, sorted keys, and functools.