Decorators In Python

Introduction

Imagine you've a cake machine that produces really nice cakes. Now, all of a sudden your customers want you to decorate your cakes with a cherry. Fair demand but if I was you, I wouldn't want to disassemble my cake machine to add this feature of putting cherries on the top.

This is not just because I'm too lazy to revisit the internals of the machine (at least this is not the sole reason) but because doing this will violate many engineering principles:

• Coupling — I shouldn't do things that make the system tightly coupled to each other
• Single responsibility — The cake machine has single responsibility; and it should continue to be its sole responsibility
• Flexibility — If the customers will ask for berries instead of cherries, I will have to do it all over.

We can address all these issues if we just make one machine, which will just decorate the cake produced by the cake machine with a cherry. Let's call it Decorator machine.

It doesn't only solve aforementioned problems but now you've a portable machine which you can use for decorating other products with cherries as well.

Decorator is a design pattern in Python that reduces code duplication and increases decoupling.

Usage

I've been working on a system that has a lot of compute intensive machine learning functions:

@cpuseconds
def train():
    """
    function to train a model
    """
    …

@cpuseconds
def predict(data, model):
    """
    function to predict using a model
    """
    …

We want to see how much of CPU seconds (code execution time on the machine) are being consumed by each customer and on what functionality.

As we already saw in our cake-machine example, editing these serene functions is not a good idea. So let us make a decorator that does that.

import time

def cpuseconds(func):
    """
    decorator method
    to calculate time elapsed
    in execution of func
    """
    def wrapper(*args, **kwargs):

        start_time = time.time()
        response = func(*args, **kwargs)
        time_elapsed = time.time()-start_time

        with open('resource-util.txt','a+') as f:
            f.write(f"{func.__name__},{time_elapsed}\n")

        return response

    return wrapper

Now, we can just decorate the existing functions with this decorator functions:

@cpuseconds
def train():
    …

@cpuseconds
def predict(data, model):
    …

When you run these functions, you'll have a file with records about time utilisation of each function call.

Note that you've changed the behaviour of the function without modifying anything inside the original function. Plan and simple.

However, there's more to this story. And before we proceed further, lets print the doc string of our original functions using dunder method __doc__.

Oops. This isn't what we expected. To understand why, we need to dive a little deeper and see how decorator has been implemented in Python.

How it's made

Decorators work on the concept of first class functions. It just means that functions in Python are treated like any other object. You can pass functions as arguments of other functions like you pass strings, integers, and any other Python object.

In the above example, by using cpuseconds decorator we're actually passing our "train" function in the function cpuseconds that actually returns another function with some added functionalities. It's almost identical but…with a cherry on the top.

So it means that @cpuseconds on top of train translates to:

train = cpuseconds(train)

We are pointing the train variable to a new function. Note that the output of the new function is a function, which hasn't been executed yet.

This means we're not actually calling our original train function, since the variable train has been reassigned to something else — something that coincidentally happens to possess some characteristics of original train function (e.g. returning the same response) but not the train function itself.

Hence, it's not surprising to see that change in train.__doc__ or even train.__name__.

@functools.wraps

Fortunately, python has a built in decorator called @functools.wraps. You can use this decorator to decorate the wrapper function. This will copy certain attributes (like __doc__, __name__, etc.) of the original function to the wrapper function, thereby preserving the info.

from functools import wraps
def cpuseconds(func):

    @wraps
    def wrapper(*args, **kwargs):
        …

    return wrapper

Now let's try again:

@wraps also comes in handy when chaining multiple decorators.

Conclusion

This was in-depth article on workings of a decorator and its real world use cases. There are other topics we didn't touch in this article e.g., class based decorators, arguments inside decorator functions, etc.

Further Readings

• Why do we need wrapper function in decorators? — Stack Overflow