Pieces of Py #1: Decorator with arguments

Posted on Tue 09 July 2019 in Python • 3 min read

This post is a simple example which explains how to create a decorator function that accepts and uses argument(s). A decorator that would accept an argument could look like this:

@mydecorator(my_argument)
def my_decorated_function(input):
    return input

A simple scenario could be that you would like to print something before and after the output of your function.

@mydecorator('-------------------')
def my_decorated_function(input):
    return input

print(my_decorated_function('Hello, World!')

In this case we would want the outcome to be:

-------------------

Hello, World!

-------------------

A decorator without argument(s)

Let's first look at a decorator without input arguments as a reference. It's usually done by creating a decorator function, with a function as an input argument. So the function you decorate will be the input to the decorator function.

Then a wrapper function is created within the decorator function that does some work. The decorator function then returns the wrapper function with some modified functionality as in the following example. Here I've put the incoming function into the result variable, and returning an f-string with that variable included.

def my_decorator(function):
    def wrapper(*args, **kwargs):
        result = function(*args, **kwargs)
        return f'Adding some text in the decorator\n{result}'
    return wrapper

So, when decorating a function with the my_decorator function, we'll see that the output of the decorated function is changed.

@my_decorator
def my_decorated_function(input):
    return input

print(my_decorated_function('Hello, World!'))

This produces the following output:

Adding some text in the decorator
Hello, World!

A decorator with an argument

But what if we don't want to hard code, in this case the text that is added to the output, and rather want to control that by adding an argument to the decorator? Well, for me it was very hard to understand how to do that at first, and it took me some googling and brainwork to get the hang of it all.

We already have a function within a function to do what we already did in the previous example. Now we also need a function, within a function, within a function. Coming from languages as C# and PowerShell, that I'd never seen before. Maybe it exists in those languages as well, but it is all new to me. So the decorator function would look like this.

def my_decorator(input_arg):

    def the_real_decorator(function):
        def wrapper(*args, **kwargs):
            result = function(*args, **kwargs)
            return f'{input_arg}\n{result}\n{input_arg}'
        return wrapper

    return the_real_decorator

And now we can add an argument to the decorator, which will be used within the decorator function. So the outer decorator function will take the argument from the decorator, and the next inner function will take the decorated function as input (as seen in the first example) and be able to use the input from the outer function.

First we return the wrapper function to the_real_decorator, and then we return the_real_decorator to the actual decorator function.

For me this was/is kind of mind blowing. But since we now use the decorator with an argument, we can see that the argument works.

@my_decorator('-------------')
def my_decorated_function(input):
    return input

print(my_decorated_function('Hello, World!'))

-------------
Hello, World!
-------------

Conclusion

I hope that in this article I've managed to explain in a simple way, how you can use argument(s) with your decorator functions in Python.

Resources

Let me know on Twitter if I can improve this article, or if you have other resources to help out with understanding this topic.

Python Pieces of Py Decorator