How to Implement a Custom Decorator That Accepts Arguments in Python

Decorators in Python are a powerful tool used to modify the behavior of functions or classes. But what if you want your decorator to accept its own arguments? This is where things can seem tricky at first, but with the right breakdown, it’s actually pretty straightforward.

In this article, we’ll explain how to implement a custom decorator that accepts arguments in Python, using clean and simple examples to help you understand and apply it in your own projects.

A decorator in Python is essentially a function that takes another function as input and returns a modified or wrapped version of that function.

Example of a simple decorator:

def simple_decorator(func):
    def wrapper():
        print("Before the function runs")
        func()
        print("After the function runs")
    return wrapper

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

say_hello()

But what if you want to pass arguments to the decorator itself, like this?

@repeat_times(3)
def say_hi():
    print("Hi!")

This is where we implement a custom decorator that accepts arguments in Python.

How to Implement a Custom Decorator That Accepts Arguments in Python

To achieve this, we use a three-layered function structure:

1. The outer function accepts the arguments for the decorator.
2. The middle function is the actual decorator that takes the target function.
3. The inner function wraps and executes the logic before/after the function call.

Let’s break it down with an example.

Example: A Decorator That Repeats a Function Multiple Times

def repeat_times(num_times):
    def decorator(func):
        def wrapper(*args, **kwargs):
            for _ in range(num_times):
                func(*args, **kwargs)
        return wrapper
    return decorator

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

greet("Alice")

What’s happening here?

• repeat_times(3) returns the decorator function.
• The decorator function takes greet as an argument.
• The wrapper function calls greet three times.

This is a textbook example of how you implement a custom decorator that accepts arguments in Python.

Example: A Decorator That Logs Messages with Custom Tags

Let’s say you want to log messages with a tag to identify their source.

def log_with_tag(tag):
    def decorator(func):
        def wrapper(*args, **kwargs):
            print(f"[{tag}] Starting '{func.__name__}'")
            result = func(*args, **kwargs)
            print(f"[{tag}] Finished '{func.__name__}'")
            return result
        return wrapper
    return decorator

@log_with_tag("DEBUG")
def process_data():
    print("Processing data...")

process_data()

You’ll get output like:

[DEBUG] Starting 'process_data'
Processing data...
[DEBUG] Finished 'process_data'

Again, we’ve used the three-layer decorator pattern to implement a custom decorator that accepts arguments in Python in a real-world use case.