Skip to main content

Decorators

In this chapter, we'll learn about decorators - a powerful Python feature for modifying function behavior. We'll cover:

  • What are decorators?
  • Creating function decorators
  • Decorators with arguments
  • Class decorators
  • @functools.wraps for preserving function metadata
  • Common built-in decorators
  • Testing decorated functions

What is a decorator?​

A decorator is a function that takes another function and extends its behavior without modifying it:

@decorator
def my_function():
    pass

# Equivalent to:
def my_function():
    pass
my_function = decorator(my_function)

Your first decorator​

Write the test first​

Let's create a decorator that logs function calls:

from decorators import log_calls

logged_calls = []


@log_calls
def greet(name):
    return f"Hello, {name}!"


def test_log_calls():
    result = greet("Alice")
    
    assert result == "Hello, Alice!"
    assert "greet" in logged_calls[-1]
    assert "Alice" in logged_calls[-1]

Implementation​

def log_calls(func):
    def wrapper(*args, **kwargs):
        call_info = f"Called {func.__name__} with args={args}, kwargs={kwargs}"
        logged_calls.append(call_info)
        return func(*args, **kwargs)
    return wrapper

How it works:

  1. log_calls receives the original function
  2. Returns a wrapper that logs then calls the original
  3. *args, **kwargs pass any arguments through

Using functools.wraps​

Without @functools.wraps, decorated functions lose their metadata:

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


@my_decorator
def greet(name):
    """Greet someone by name."""
    return f"Hello, {name}!"


print(greet.__name__)  # "wrapper" - wrong!
print(greet.__doc__)   # None - wrong!

Fix with @functools.wraps:

from functools import wraps


def my_decorator(func):
    @wraps(func)
    def wrapper(*args, **kwargs):
        return func(*args, **kwargs)
    return wrapper


@my_decorator
def greet(name):
    """Greet someone by name."""
    return f"Hello, {name}!"


print(greet.__name__)  # "greet" - correct!
print(greet.__doc__)   # "Greet someone by name." - correct!

Testing metadata preservation​

def test_decorator_preserves_metadata():
    @log_calls
    def example_func():
        """Example docstring."""
        pass
    
    assert example_func.__name__ == "example_func"
    assert example_func.__doc__ == "Example docstring."

Practical decorators​

Timing decorator​

import time
from functools import wraps


def timing(func):
    @wraps(func)
    def wrapper(*args, **kwargs):
        start = time.time()
        result = func(*args, **kwargs)
        duration = time.time() - start
        print(f"{func.__name__} took {duration:.4f}s")
        return result
    return wrapper

Test:

def test_timing_decorator():
    @timing
    def slow_function():
        time.sleep(0.1)
        return "done"
    
    result = slow_function()
    
    assert result == "done"

Retry decorator​

import time
from functools import wraps


def retry(max_attempts=3, delay=1):
    def decorator(func):
        @wraps(func)
        def wrapper(*args, **kwargs):
            last_exception = None
            for attempt in range(max_attempts):
                try:
                    return func(*args, **kwargs)
                except Exception as e:
                    last_exception = e
                    if attempt < max_attempts - 1:
                        time.sleep(delay)
            raise last_exception
        return wrapper
    return decorator

Test:

def test_retry_succeeds_eventually():
    attempts = []
    
    @retry(max_attempts=3, delay=0)
    def flaky_function():
        attempts.append(1)
        if len(attempts) < 3:
            raise ValueError("Not yet!")
        return "success"
    
    result = flaky_function()
    
    assert result == "success"
    assert len(attempts) == 3

Decorators with arguments​

Decorators that accept arguments need an extra layer:

def repeat(times):
    def decorator(func):
        @wraps(func)
        def wrapper(*args, **kwargs):
            result = None
            for _ in range(times):
                result = func(*args, **kwargs)
            return result
        return wrapper
    return decorator


@repeat(times=3)
def say_hello():
    print("Hello!")
    return "done"

Testing parameterized decorators​

def test_repeat_decorator():
    calls = []
    
    @repeat(times=3)
    def track_calls():
        calls.append(1)
        return "ok"
    
    result = track_calls()
    
    assert result == "ok"
    assert len(calls) == 3

Validation decorator​

A common use case - validating function arguments:

from functools import wraps


def validate_positive(func):
    @wraps(func)
    def wrapper(*args, **kwargs):
        for arg in args:
            if isinstance(arg, (int, float)) and arg < 0:
                raise ValueError(f"Arguments must be positive, got {arg}")
        return func(*args, **kwargs)
    return wrapper


@validate_positive
def square_root(n):
    return n ** 0.5

Test:

import pytest


def test_validate_positive():
    @validate_positive
    def double(n):
        return n * 2
    
    assert double(5) == 10


def test_validate_positive_rejects_negative():
    @validate_positive
    def double(n):
        return n * 2
    
    with pytest.raises(ValueError, match="must be positive"):
        double(-5)

Caching decorator​

The @functools.lru_cache decorator caches results:

from functools import lru_cache


@lru_cache(maxsize=128)
def fibonacci(n):
    if n < 2:
        return n
    return fibonacci(n - 1) + fibonacci(n - 2)


# Without cache: O(2^n)
# With cache: O(n)

Test:

def test_cached_fibonacci():
    result = fibonacci(30)
    assert result == 832040
    
    # Verify caching
    info = fibonacci.cache_info()
    assert info.hits > 0

Class decorators​

Decorators can modify classes:

def singleton(cls):
    instances = {}
    
    @wraps(cls)
    def get_instance(*args, **kwargs):
        if cls not in instances:
            instances[cls] = cls(*args, **kwargs)
        return instances[cls]
    
    return get_instance


@singleton
class Database:
    def __init__(self):
        self.connected = True

Test:

def test_singleton():
    db1 = Database()
    db2 = Database()
    
    assert db1 is db2

Adding methods with class decorators​

def add_repr(cls):
    def __repr__(self):
        attrs = ", ".join(
            f"{k}={v!r}"
            for k, v in self.__dict__.items()
        )
        return f"{cls.__name__}({attrs})"
    
    cls.__repr__ = __repr__
    return cls


@add_repr
class Point:
    def __init__(self, x, y):
        self.x = x
        self.y = y

Test:

def test_add_repr():
    point = Point(3, 4)
    assert repr(point) == "Point(x=3, y=4)"

Common built-in decorators​

@property​

class Circle:
    def __init__(self, radius):
        self._radius = radius
    
    @property
    def radius(self):
        return self._radius
    
    @radius.setter
    def radius(self, value):
        if value < 0:
            raise ValueError("Radius must be positive")
        self._radius = value

@staticmethod and @classmethod​

class Temperature:
    def __init__(self, celsius):
        self.celsius = celsius
    
    @classmethod
    def from_fahrenheit(cls, f):
        return cls((f - 32) * 5 / 9)
    
    @staticmethod
    def is_freezing(celsius):
        return celsius <= 0

@dataclass​

from dataclasses import dataclass


@dataclass
class Point:
    x: float
    y: float
    
    def distance_from_origin(self):
        return (self.x ** 2 + self.y ** 2) ** 0.5

Stacking decorators​

Multiple decorators apply bottom-up:

@decorator1
@decorator2
@decorator3
def my_function():
    pass

# Equivalent to:
my_function = decorator1(decorator2(decorator3(my_function)))

Example:

@timing
@retry(max_attempts=3)
def fetch_data():
    # Retry is applied first, then timing
    pass

Testing decorated functions​

Testing the wrapper directly​

def test_log_calls_decorator():
    calls = []
    
    def capture_log(func):
        @wraps(func)
        def wrapper(*args, **kwargs):
            calls.append((func.__name__, args, kwargs))
            return func(*args, **kwargs)
        return wrapper
    
    @capture_log
    def add(a, b):
        return a + b
    
    result = add(2, 3)
    
    assert result == 5
    assert calls == [("add", (2, 3), {})]

Testing decorated behavior​

def test_retry_behavior():
    attempt_count = 0
    
    @retry(max_attempts=3, delay=0)
    def failing_function():
        nonlocal attempt_count
        attempt_count += 1
        raise ValueError("Always fails")
    
    with pytest.raises(ValueError):
        failing_function()
    
    assert attempt_count == 3

Wrapping up​

We've covered:

  • Basic decorators - Functions that modify other functions
  • @functools.wraps - Preserve function metadata
  • Decorators with arguments - Add configuration to decorators
  • Class decorators - Modify classes
  • Built-in decorators - @property, @staticmethod, @classmethod, @dataclass
  • Stacking decorators - Apply multiple decorators

Key takeaways​

  1. Always use @functools.wraps in your decorators
  2. Decorators with arguments need three levels of nesting
  3. Test both the decorator behavior and metadata preservation
  4. Stack decorators when you need multiple behaviors
  5. Use built-in decorators for common patterns

Decorators are a powerful tool for adding cross-cutting concerns like logging, caching, validation, and retry logic without cluttering your core business logic.