Writing Decorators

Welcome to your deep dive into writing decorators in Python — a powerful and elegant feature that lets you modify or enhance the behavior of functions or methods without changing their actual code. Whether you’ve used decorators like @staticmethod or @property before, writing your own custom decorators can unlock a new level of expressive and reusable code.

In this lesson, we’ll explore what decorators are, how they work under the hood, and guide you step-by-step through writing your own. Along the way, you’ll encounter practical examples and learn best practices to avoid common pitfalls. By the end, you’ll confidently create decorators that suit your programming needs.

Understanding What a Decorator Is

At its core, a decorator is a function that takes another function as input and returns a new function that adds some kind of enhancement or modification to the original. This means decorators are higher-order functions — functions that operate on other functions.

Imagine you have a simple function that prints a greeting:

def greet():
    print("Hello, world!")

greet()

Now, suppose you want to enhance this function by adding a message before and after the greeting. You could modify greet() directly, but what if you want to keep greet() clean and reusable, and add this extra behavior only in certain contexts? This is where decorators shine.

How Decorators Work: Functions Returning Functions

To understand decorators, it's essential to grasp that functions in Python are first-class objects. You can pass them around as arguments, return them from other functions, and assign them to variables.

Here’s a simple function that returns another function:

def outer():
    def inner():
        print("I am inside inner()")
    return inner

f = outer()
f()

Here, the outer() function returns the inner() function without running it. We then assign f to this returned function and call it. This is the pattern decorators use.

Writing Your First Decorator

Let's write a simple decorator that adds a line before and after calling any function it decorates. We'll call it announce.

def announce(func):
    def wrapper():
        print("About to run the function...")
        func()
        print("Function has finished running.")
    return wrapper

@announce
def say_hello():
    print("Hello!")

say_hello()

Here’s what happens step-by-step:

• announce is a function that takes func as an argument.
• Inside, it defines a nested function wrapper(), which adds extra behavior before and after calling func().
• It returns this wrapper function.
• The @announce syntax above say_hello is syntactic sugar for say_hello = announce(say_hello).
• Calling say_hello() now actually calls the wrapper() function, which adds the announcements.

Supporting Arguments: Making Decorators Flexible

What if your original function takes arguments? The decorator’s wrapper function must be able to accept and pass these arguments along. To do this, use *args and **kwargs to capture any number of positional and keyword arguments.

def announce(func):
    def wrapper(*args, **kwargs):
        print(f"Calling {func.__name__}...")
        result = func(*args, **kwargs)
        print(f"{func.__name__} finished.")
        return result
    return wrapper

@announce
def greet(name):
    print(f"Hello, {name}!")

greet("Alice")

This pattern, def wrapper(*args, **kwargs):, is the standard way to write decorators that support any function signature.

Preserving Function Metadata with functools.wraps

One downside of writing decorators is that the decorated function loses some useful metadata like its name and docstring, because the wrapper replaces the original function.

To fix this, Python provides functools.wraps, a decorator for your wrapper function. It copies the metadata from the original function to the wrapper, preserving important attributes.

from functools import wraps

def announce(func):
    @wraps(func)
    def wrapper(*args, **kwargs):
        print(f"Calling {func.__name__}...")
        result = func(*args, **kwargs)
        print(f"{func.__name__} finished.")
        return result
    return wrapper

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

print(greet.__name__)
print(greet.__doc__)

Without @wraps(func), greet.__name__ would return wrapper instead of greet, and the docstring would be lost.

Decorators with Arguments: Making Them More Powerful

Sometimes, you want your decorator itself to accept arguments. For example, a decorator that repeats a function multiple times — you might want to specify how many times.

To do this, you add another outer layer of function, which takes the decorator arguments, and returns the actual decorator.

from functools import wraps

def repeat(num_times):
    def decorator_repeat(func):
        @wraps(func)
        def wrapper(*args, **kwargs):
            for _ in range(num_times):
                func(*args, **kwargs)
        return wrapper
    return decorator_repeat

@repeat(3)
def say_hello():
    print("Hello!")

say_hello()

Here’s the flow:

• repeat takes num_times as an argument and returns decorator_repeat.
• decorator_repeat is the actual decorator that takes the function.
• The wrapper calls the original function the specified number of times.
• Applying @repeat(3) means calling repeat(3), which returns a decorator that is then applied.

Common Use Cases for Writing Your Own Decorators

When writing decorators, think about the cross-cutting concerns that you want to reuse and isolate cleanly:

• Logging: Automatically log function calls and arguments.
• Timing: Measure the execution time of functions.
• Access control: Check user permissions before allowing a function to run.
• Caching: Store results of expensive function calls to avoid repetition.
• Retry logic: Automatically retry a function on failure.

Writing decorators for these cases helps keep your core logic clean and your code DRY (Don’t Repeat Yourself).

Example: A Timer Decorator

Let’s write a decorator that prints how long a function takes to run. This is a classic example that uses the time module.

import time
from functools import wraps

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

@timer
def waste_time(num):
    total = 0
    for i in range(num):
        total += i*i
        time.sleep(0.01)
    return total

waste_time(5)

What Happens When You Stack Decorators?

You can apply multiple decorators on the same function. They are applied from the bottom up, meaning the decorator closest to the function is applied first.

@decorator_one
@decorator_two
def func():
    pass

func → decorator_two(func) → decorator_one(decorator_two(func))

Each decorator wraps the function returned by the previous decorator, creating a layered effect.

Best Practices When Writing Decorators

• Use functools.wraps to preserve function metadata.
• Support arbitrary arguments using *args and **kwargs to make your decorator generic.
• Keep decorators simple and focused on one responsibility.
• Test decorated functions to ensure behavior matches expectations.
• Document your decorators clearly so users understand their effect.

Wrapping Up

Writing decorators in Python gives you a powerful tool to extend and modify function behavior cleanly. By mastering the pattern of functions returning wrapper functions, supporting arguments, and preserving metadata, you’ll be able to craft decorators for a wide range of use cases.

Try writing your own decorator now that logs function arguments or retries a function on exception. Experimenting is the best way to internalize these concepts and see the magic of decorators in action.