Python Decorators: A Comprehensive Guide

Photo by Element5 Digital on Unsplash

Decorator is a common concept in Python and you will see them everywhere. It can be a bit hard to understand at first (I was there) but it’s actually a very simple concept. Decorators give a cleaner way to modify or extend the behavior of a function without needing to change it’s source code. Confused? Don’t worry, in this article, I will try to demystify decorators for you.

Let’s try to understand decorators by exploring a problem first, so that we have a clear idea why and when to use a decorator.

Suppose we have the following utility function in our code base,

# utility.py
def n_squares(n):
  squares = []
  for x in range(1, n+1):
    squares.append(x**2)
  return squares

It simply returns squares of 1 to n in a list. And of course, we have probably used this function in many places in our code base, like the following,

# service.py
from utility import n_squares

fifty_squares = n_squares(50)
five_hundred_squares = n_squares(500)
five_thousand_squares = n_squares(5000)

Now, suddenly an additional requirement came in. You have to calculate the execution time of this utility function every time it is called and log it.

Say, your first thought was to solve it this way,

# service.py
from utility import n_squares
from time import time

start = time()
fifty_squares = n_squares(50)
end = time()
print('execution time:', end - start)

start = time()
five_hundred_squares = n_squares(500)
end = time()
print('execution time:', end - start)

start = time()
five_thousand_squares = n_squares(5000)
end = time()
print('execution time:', end - start)

This is the ugliest solution. You have to repeat these lines everywhere in your code where you have called the n_squares function and will call in the future. Not recommended.

Say, this is your second approach,

# utility.py
from time import time

def n_squares(n):
  start = time()
  squares = []
  for x in range(1, n+1):
    squares.append(x**2)
  end = time()
  print('execution time:', end - start)
  return squares

Here, you have updated the code inside the function directly, which is cleaner than the previous solution. But, what if, this is not the only function where you need to log the execution time. What if you have other functions in your utility file and your senior engineer has asked you to log the execution time for every other functions in the utility file? You will need to repeat the extra three lines of code in every single utility function, which is bad. It breaks the DRY (Don’t Repeat Yourself) principle. If the added logic was something much more complex, or if you later need to change them, you will have to go through all the functions again, and repeat the process and might make a mistake somewhere.

So, how can we DRY this? Let’s move the execution time calculation logic to a separate function.

# utility.py
from time import time

def log_execution_time(fn):
  def wrapped_function(*args, **kwargs):
    start = time()
    res = fn(*args, **kwargs)
    end = time()
    print('execution time:', end - start)
    return res
  return wrapped_function

(*args and **kwargs means, ‘all the positional and keyword arguments that were passed’. It’s a way to keep the function more generic by not limiting it to any specific number of arguments.)

Here we defined a function named log_execution_time which returns another function. What does that mean? It means, when you call this function and assign the returned value to a variable, that variable will behave like a function (the function that was returned). See the following example if you want to understand the concept.

def top_function():
  def to_be_returned_function():
    print('I am from the inner function!')

  return to_be_returned_function

my_variable = top_function() # prints nothing

# now you can treat my_variable like a function, because it IS a function
my_variable() # prints -> I am from the inner function!
              # because now you are invoking the 'to_be_returned_function'

Our log_execution_time function takes a function for which we want to print the execution time. It returns a function which wraps the given fn with the execution time calculation logic, prints it, and then returns the returned value of fn.

Let’s continue in our utility file. We can now use the log_execution_time function to update the behavior of the n_squares function.

# utility.py
from time import time

def log_execution_time(fn):
  def wrapped_function(*args, **kwargs):
    start = time()
    res = fn(*args, **kwargs)
    end = time()
    print('execution time:', end - start)
    return res
  return wrapped_function

def n_squares(n):
  squares = []
  for x in range(1, n+1):
    squares.append(x**2)
  return squares

n_squares = log_execution_time(n_squares)

We kept the original definition of the n_squares function. In python, the function name we use to define the functions is like any other objects, so we can pass them to other function, and also reassign them! And that’s what we did in the last line. We just reassigned n_squares to the function that is returned from the log_execution_time function.

Notice that, for this, we needed to use the word n_squares three times, first when we were defining the original function, second when we wanted to reassign its value, and third when we passed it to log_execution_time function as the parameter. This can be error-prone if we have functions with almost same name. For example, can you spot the bug in the following code block?

def n_squares(n):
  squares = []
  for x in range(1, n+1):
    squares.append(x**2)
  return squares

def n_square(n):
  return n**2

n_squares = log_execution_time(n_square)

We passed n_square instead of n_squares (with ‘s’ that returns a list). So now the n_squares function does a completely different thing! This can happen. So python gives us an easier syntax to save us from this type of mistake. Decorator!

Instead of calling the log_execution_time function passing the function name we want to modify and then storing the returned function to the original function name, we can simply put the log_execution_time function with a “@” symbol at the top of the function definition. Python will handle the rest.

# utility.py
from time import time

def log_execution_time(fn):
  def wrapped_function(*args, **kwargs):
    start = time()
    res = fn(*args, **kwargs)
    end = time()
    print('execution time:', end - start)
    return res
  return wrapped_function

@log_execution_time  # <-- decorator in play here!
def n_squares(n):
  squares = []
  for x in range(1, n+1):
    squares.append(x**2)
  return squares

And that’s it! That’s Decorator. It’s just a regular function, that takes another function and adds additional logic to it. We can use decorator by using the @ symbol and it will take the function below it as it’s parameter and reassign it accordingly.

What if I need to pass parameter?

If you need to pass parameters to your decorator, just wrap the whole thing in another function! Let’s look at another example. Say we want to add a decorator that increases the returned value of a function by an specified percentage

def return_with_increased_value(fn):
  rate = 0.15
  def wrapped(*args, **kwargs):
    res = fn(*args, **kwargs)
    return res + (res * rate)
  return wrapped

We have hard-coded the rate. If you wanted to make this dynamic, this is how you can do it.

def with_percentage(rate):
  def return_with_increased_value(fn):
    def wrapped(*args, **kwargs):
      res = fn(*args, **kwargs)
      return res + (res * rate)
    return wrapped
  return return_with_increased_value

And this is how you can use it,

@with_percentage(0.15)
def calculate_price(items):
  pass

To simplify what is happening, the call to with_percentage returns the return_with_increased_value function, and then that function receives calculate_price and modifies it. In non-decorator syntax, this is what’s happening,

decorator = with_percentage(0.15)
calculate_price = decorator(calculate_price)

Passing parameter to function from decorator

You can use decorator to pass parameter value to your function. Suppose, you are writing a method to add a new entry in a database. Your function needs a session object that handles the db operations.

def create_entry(session, data):
  pass

You want the caller of this function to only pass the data, and not worry about the session object (maybe because there are some repeated steps that needs to happen every time, and you want to keep the code DRY). This is what you can do,

from somewhere import db

def use_session(fn):
  def wrapped(*args, **kwargs):
    session = db.get_session()
    res = fn(session, *args, **kwargs
    # the repeated logics that you don't need to repeat anymore in your code
    session.commit()
    session.close()
    return res
  return wrapped

@use_session
def create_entry(session, data):
  pass

Now wherever you need to call your create_entry function, you don’t need to pass the session object, and whatever function now needs to use the session object, they don’t need to repeat the same steps. The decorator will take care of it. You will need to pass the other parameter(s) only, like this,

data = {"firstname": "John", "lastname": "Doe"}
user = create_entry(data)

Can I use a class as a decorator?

(This part is bit advanced)

We have defined the decorators as functions. But can decorators be something else, like a Class?

In python, syntactically there is no difference between calling a function and a class. So what is stopping us from using a class?

class MyClass:
  pass

def my_function:
  pass

@Myclass
def function_a():
  pass

# translates to -> function_a = MyClass(function_a)

@my_function
def function_b():
  pass

# translates to -> function_b = my_function(function_b)

The decorator functions we define, return a function, which then gets called when we use the decorated function (functon_a in the above example). But calling a Class constructor does not return a function, it returns an instance of that class. So what will happen if you call function_b? Can you still use function_b like a function?

Yes you can! We have a dunder method __call__ for that. The __call__ method is used to make the instance of a class callable. We can utilize this to make a class based decorator.

Class MyClass:
  def __init__(self, fn):
    self.fn = fn

  def __call__(self, *args, **kwargs):
    # Do: logic before calling the original function
    res = self.fn(*args, **kwargs)
    # Do: logic after calling the original function
    return res

That’s all you need to know about decorators. Happy coding!

(If you found this article helpful, give it some claps! If you want to connect with me, send me a request on Twitter@AhsanShihab_.)