Decorators: Wrapping Functions

The idea, without any syntax

def announce(func):
    def wrapper():
        print("before")
        result = func()
        print("after")
        return result
    return wrapper


def greet():
    print("hello")


greet = announce(greet)      # rebind the name to the wrapper
greet()

before
hello
after

Three facts make this work, all covered earlier in this path: functions are objects, a function can return a function, and the inner function keeps hold of func through a closure.

The @ syntax

@announce
def greet():
    print("hello")


greet()

@announce above a def means exactly greet = announce(greet), run immediately after the function is defined. That is the entire feature. Everything else is detail.

Handling any signature

def announce(func):
    def wrapper(*args, **kwargs):          # accept anything
        print(f"calling {func.__name__}")
        result = func(*args, **kwargs)     # pass it all on
        print(f"  returned {result!r}")
        return result                      # give the caller the result
    return wrapper


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


@announce
def greet(name, greeting="Hello"):
    return f"{greeting}, {name}"


print(add(2, 3))
print(greet("Meera", greeting="Welcome"))

Three rules for every wrapper: accept *args, **kwargs, pass them through unchanged, and return the result. Forgetting the return is the most common decorator bug - the decorated function silently starts returning None.

functools.wraps

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


@announce
def greet(name):
    """Return a greeting."""
    return f"Hello, {name}"


print(greet.__name__)      # wrapper       <- the identity is gone
print(greet.__doc__)       # None

from functools import wraps


def announce(func):
    @wraps(func)                            # copy name, docstring, signature
    def wrapper(*args, **kwargs):
        return func(*args, **kwargs)
    return wrapper


@announce
def greet(name):
    """Return a greeting."""
    return f"Hello, {name}"


print(greet.__name__)         # greet
print(greet.__doc__)          # Return a greeting.
print(greet.__wrapped__)      # the original function, still reachable

Without @wraps, every decorated function in your program is called wrapper. That breaks documentation tools, test discovery, debuggers and tracebacks. Use it always.

Useful decorators

Timing

import time
from functools import wraps


def timed(func):
    @wraps(func)
    def wrapper(*args, **kwargs):
        start = time.perf_counter()
        try:
            return func(*args, **kwargs)
        finally:
            elapsed = time.perf_counter() - start
            print(f"{func.__name__} took {elapsed:.4f}s")
    return wrapper


@timed
def slow_sum(n):
    return sum(range(n))


slow_sum(1_000_000)

The finally means the timing is reported even when the function raises.

Caching

from functools import wraps


def memoise(func):
    cache = {}

    @wraps(func)
    def wrapper(*args):
        if args not in cache:
            cache[args] = func(*args)
        return cache[args]

    wrapper.cache = cache                 # expose it for inspection
    wrapper.clear = cache.clear
    return wrapper


@memoise
def fib(n):
    return n if n < 2 else fib(n - 1) + fib(n - 2)


print(fib(50))
print(len(fib.cache), "entries")

In real code use functools.cache, which does this properly. Writing it once by hand is the clearest way to see that a decorator is just a closure.

Validation

from functools import wraps


def positive_arguments(func):
    @wraps(func)
    def wrapper(*args, **kwargs):
        for value in list(args) + list(kwargs.values()):
            if isinstance(value, (int, float)) and value <= 0:
                raise ValueError(f"{func.__name__}: {value} must be positive")
        return func(*args, **kwargs)
    return wrapper


@positive_arguments
def rectangle_area(width, height):
    return width * height


print(rectangle_area(3, 4))
try:
    rectangle_area(3, -4)
except ValueError as error:
    print(error)

Retrying

import time
from functools import wraps


def retry(attempts=3, delay=0.1, exceptions=(Exception,)):
    def decorator(func):
        @wraps(func)
        def wrapper(*args, **kwargs):
            last = None
            for attempt in range(1, attempts + 1):
                try:
                    return func(*args, **kwargs)
                except exceptions as error:
                    last = error
                    print(f"  attempt {attempt} failed: {error}")
                    if attempt < attempts:
                        time.sleep(delay)
            raise last
        return wrapper
    return decorator


calls = 0


@retry(attempts=3, delay=0.01, exceptions=(ConnectionError,))
def flaky():
    global calls
    calls += 1
    if calls < 3:
        raise ConnectionError("network unavailable")
    return "connected"


print(flaky())

Decorators that take arguments

A decorator with arguments needs one more layer. Read the nesting from the outside in:

@repeat(3)
def greet(): ...

is       greet = repeat(3)(greet)
          |        |        |
          |        |        +-- the decorator is applied
          |        +----------- repeat(3) RETURNS a decorator
          +-------------------- the final name is bound to the wrapper

from functools import wraps


def repeat(times):
    def decorator(func):
        @wraps(func)
        def wrapper(*args, **kwargs):
            results = []
            for _ in range(times):
                results.append(func(*args, **kwargs))
            return results
        return wrapper
    return decorator


@repeat(3)
def roll():
    import random
    return random.randint(1, 6)


print(roll())

Working with and without arguments

from functools import wraps


def logged(func=None, *, prefix="LOG"):
    def decorator(inner):
        @wraps(inner)
        def wrapper(*args, **kwargs):
            print(f"[{prefix}] {inner.__name__}")
            return inner(*args, **kwargs)
        return wrapper

    if func is None:            # called as @logged(prefix="X")
        return decorator
    return decorator(func)      # called as @logged


@logged
def a():
    return 1


@logged(prefix="AUDIT")
def b():
    return 2


a()
b()

Stacking decorators

from functools import wraps


def bold(func):
    @wraps(func)
    def wrapper(*args, **kwargs):
        return f"<b>{func(*args, **kwargs)}</b>"
    return wrapper


def italic(func):
    @wraps(func)
    def wrapper(*args, **kwargs):
        return f"<i>{func(*args, **kwargs)}</i>"
    return wrapper


@bold
@italic
def text():
    return "hello"


print(text())         # <b><i>hello</i></b>

Decorators apply bottom up. @italic wraps the function first, then @bold wraps that. The one nearest the def is applied first, and so ends up innermost.

Decorating methods

from functools import wraps


def audit(func):
    @wraps(func)
    def wrapper(self, *args, **kwargs):      # self arrives as the first argument
        print(f"{type(self).__name__}.{func.__name__} called")
        return func(self, *args, **kwargs)
    return wrapper


class Account:
    def __init__(self, balance=0):
        self.balance = balance

    @audit
    def deposit(self, amount):
        self.balance += amount
        return self.balance


a = Account()
print(a.deposit(100))

A generic *args wrapper handles methods without change, because self is simply the first positional argument.

Order matters with property and classmethod

class Example:
    @property                # property must be OUTERMOST
    @audit
    def value(self):
        return 42

    @classmethod             # classmethod must be OUTERMOST
    @audit
    def build(cls):
        return cls()

Class decorators

def auto_repr(cls):
    """Add a repr built from the instance dictionary."""
    def __repr__(self):
        fields = ", ".join(f"{k}={v!r}" for k, v in vars(self).items())
        return f"{cls.__name__}({fields})"

    cls.__repr__ = __repr__
    return cls


@auto_repr
class Product:
    def __init__(self, name, price):
        self.name = name
        self.price = price


print(Product("Notebook", 250))     # Product(name='Notebook', price=250)

def singleton(cls):
    """Ensure only one instance of the class ever exists."""
    instances = {}

    def get_instance(*args, **kwargs):
        if cls not in instances:
            instances[cls] = cls(*args, **kwargs)
        return instances[cls]

    return get_instance


@singleton
class Settings:
    def __init__(self):
        print("loading settings")
        self.theme = "dark"


a = Settings()      # loading settings
b = Settings()      # nothing printed
print(a is b)       # True

Decorators in the standard library

from functools import cache, wraps, singledispatch, total_ordering
from dataclasses import dataclass


@cache                  # memoise
@wraps                  # preserve identity
@property               # method as attribute
@staticmethod           # no self
@classmethod            # cls instead of self
@dataclass              # generate __init__, __repr__, __eq__
@total_ordering         # fill in comparison methods
@singledispatch         # dispatch on argument type
@abstractmethod         # must be overridden

A worked example

import time
from functools import wraps


def instrument(name=None, log_args=False):
    """Time a function, count its calls and optionally log its arguments."""

    def decorator(func):
        label = name or func.__name__

        @wraps(func)
        def wrapper(*args, **kwargs):
            wrapper.calls += 1
            if log_args:
                print(f"  {label}{args}{kwargs or ''}")
            start = time.perf_counter()
            try:
                result = func(*args, **kwargs)
            except Exception as error:
                wrapper.failures += 1
                print(f"  {label} raised {type(error).__name__}")
                raise
            finally:
                wrapper.total_time += time.perf_counter() - start
            return result

        wrapper.calls = 0
        wrapper.failures = 0
        wrapper.total_time = 0.0
        wrapper.report = lambda: (
            f"{label}: {wrapper.calls} calls, "
            f"{wrapper.failures} failures, "
            f"{wrapper.total_time:.4f}s total"
        )
        return wrapper

    return decorator


@instrument(log_args=True)
def divide(a, b):
    return a / b


divide(10, 2)
divide(9, 3)
try:
    divide(1, 0)
except ZeroDivisionError:
    pass

print(divide.report())

Common mistakes

• Forgetting to return the result from the wrapper.
• Forgetting @wraps.
• Writing @decorator() when the decorator takes no arguments, or @decorator when it does.
• Returning decorator(func) instead of decorator from a parameterised decorator.
• Getting the stacking order wrong with @property or @classmethod.
• Putting expensive work in the decorator body, which runs at import time.
• Caching a function whose result depends on outside state.

Best practices

• Always use @wraps.
• Always accept *args, **kwargs and return the result.
• Keep each decorator to one concern.
• Attach counters and controls to the wrapper so they can be inspected.
• Use try ... finally when the decorator must always run cleanup.
• Reach for a standard library decorator before writing your own.

Practice

1. Write a decorator that prints the arguments and the return value of any function.
2. Write a decorator taking a maximum duration and warning when a call exceeds it.
3. Demonstrate the difference @wraps makes to __name__ and help().
4. Stack three decorators and predict the output order before running it.
5. Write a class decorator that records every instance created.

Conclusion

@decorator means func = decorator(func). Everything else - arguments, stacking, methods, classes - is that one substitution repeated. Accept *args, **kwargs, return the result, and never omit @wraps.