Python Decorators - Deep Dive in 4 Parts (2/4)

This is a 4 parts series:

Part 1 : Intro + build your first decorator
Part 2 (current): Decorators for functions that receive parameters and return values, decorators that receive arguments, decorators in classes + advanced exercises with solutions
Part 3: Concatenating decorators (TBD)
Part 4: Property decorator (TBD)

This article is Part 2 in series about Python Decorators, and it assumes you have basic knowledge in decorators. If you don't, or if you haven't read the first article in the series, feel free to do so here.

In this article we will discuss implementing more advanced decorators, in particular:

Implementing decorators for any function, including functions that receive parameters and return values
Implementing decorators that receive arguments by themselves (i.e. modifying decorators behavior)
Implementing decorators inside classes
Exercises to improve your skills + solutions

If you prefer watching and listening, feel free to check out my video explaining these topics:

https://youtu.be/CUZqICabccw

1. Recap

Let's start from a short recap of what decorator is and what we learned in a previous part of this series. Below is an implementation of a simple greeting_decorator that prints out beautiful messages before and after running the function it decorates.

def greeting_decorator(other_func):

    def greeting_func():
        print(f"\n---------------------------------------\n"
              f"Hello!\nWelcome to function {other_func.__name__}!\n"
              f"---------------------------------------\n")
        # Note we call the function we received as parameter
        # inside our inner function
        other_func()
        print(f"\n---------------------------------------\n"
              f"Good-bye!\nThanks for using function {other_func.__name__}!"
              f"\n---------------------------------------\n")

    return greeting_func


@greeting_decorator
def sum_of_digits():
  num = int(input(f"Insert a number: "))
  digits_sum = 0
  while num != 0:
    digits_sum += num % 10
    num = num // 10
  print(f"The sum of digits is: {digits_sum}")

Since we defined sum_of_digits decorated with greeting_decorator, the actual sum_of_digits that was stored by Python interpreter is a function returned by greeting_decorator. In other words, the actual code that was stored with identifier sum_of_digits is a return value of:

greeting_decorator(sum_of_digits)

Hence, when running the following line:

sum_of_digits()

Python interpreter executes decorator code, that in turn displays welcome and goodbye messages and executes the original sum_of_digits code, so the output for this line will be as follows:

You might noticed that I chose to implement sum_of_digits in a bit weird way: instead of passing a number as argument and returning the sum of its digits as return value, I intentionally chose to get a number as an input and print the result inside the function. It has been done to simplify the implementation of our decorator. But now, after we have basic knowledge in decorators, we can move further.

2. Decorating functions that receive parameters and return values

Let's change our sum_of_digits function to a standard implementation that makes more sense. Now our function will receive a number as a parameter and return the sum.

@greeting_decorator
def sum_of_digits(num):
  digits_sum = 0
  while num != 0:
    digits_sum += num % 10
    num = num // 10
  return digits_sum

If we'll try running:

sum_of_digits(235)

with greeting_decorator left as is, we'll get an exception:

It happens because our greeting_func implemented inside greeting_decorator indeed takes no arguments, and as we already know - what will really get called when running sum_of_digits(235), is:

greeting_decorator(sum_of_digits)(235)

It means that we have to change greeting_func returned from greeting_decorator to get an argument and to return a value, as follows:

def greeting_decorator(other_func):

    # Note: now greeting_func receives a parameter
    def greeting_func(num):

        print(f"\n---------------------------------------\n"
              f"Hello!\nWelcome to function {other_func.__name__}!\n"
              f"---------------------------------------\n")

        # Return value from the other_func is stored in a temp variable
        # to be returned later, after printing the good-bye message
        result = other_func(num)

        print(f"\n---------------------------------------\n"
              f"Good-bye!\nThanks for using function {other_func.__name__}!"
              f"\n---------------------------------------\n")

        # Note: greeting_func returns a value!
        return result

    return greeting_func

Now, when calling:

ret_sum = sum_of_digits(345)
print(f"Sum of digits for 345 is: {ret_sum}")

we will get an expected output:

Though our greeting_decorator works as expected now, the way it handles function arguments is not generic enough. What will happen if we'll try to use greeting_decorator for function that takes more than one argument? Or if the decorated function will include both required and keyword arguments? Or if the function takes no arguments? Of course, it will again raise an exception for all these cases, since now we require that the decorated function takes exactly one required argument.

Happily, in Python we can easily rewrite our decorator to support all the cases described above and even more, by using args and *kwargs signature:

def greeting_decorator(other_func):

    # Note usage of *args, **kwargs
    def greeting_func(*args, **kwargs):

        print(f"\n---------------------------------------\n"
              f"Hello!\nWelcome to function {other_func.__name__}!\n"
              f"---------------------------------------\n")

        # Calling the function with all the parameters sent
        result = other_func(*args, **kwargs)

        print(f"\n---------------------------------------\n"
              f"Good-bye!\nThanks for using function {other_func.__name__}!"
              f"\n---------------------------------------\n")

        # Note: greeting_func returns a value!
        return result

    return greeting_func

Now, we can use our greeting_decorator with any function!

To summarize: when implementing a decorator, you should work according to the following template:

# Decorator template
def my_decorator(other_func):

    # Include *args, **kwargs in the signature
    def decorated_func(*args, **kwargs):

        # Optionally do some stuff before calling the original function

        # IMPORTANT! Call the original function passing it *args, **kwargs
        result = other_func(*args, **kwargs)

        # Optionally do some stuff after calling the original function

        # IMPORTANT! Don't forget to return the result of the original function
        return result

    # Return the function you created
    return decorated_func

Now it's time for an exercise. I encourage you to try solving it by yourself 💪 before looking at my solution.

Exercise 1: Implement a decorator that logs execution time of a function

Implement a decorator performance_log that prints amount of time (in seconds) that it takes to a function to complete execution. It can be very useful for debugging and profiling purposes.

Test your decorator with the functions provided below.

Hint: try using time.perf_counter() to measure time

# Implement this
def performance_log(func):
  pass


# Function to test with
@performance_log
def long_running_func(num, iters):
  val = 1
  for i in range(iters):
    val *= num
  return val

Below are example outputs that you should get when calling long_running_func after you implement the decorator:

res = long_running_func(17, 1000)
print(f"The result is: {res}")

Expected output:

Execution time is: 0.00027704099989023234
The result is 281139182902740093173255193460516433570993900889613439277903794687196783510046951084......

An optional solution for this exercise can be found here, but I'm sure you don't need it because you solved it by yourself 😅

Spoiler alert⚠️: Next section is based on the decorator performance_log

3. Passing parameters to decorators

After we implemented performance_log decorator that logs function execution time in seconds, we want to add more flexibility to it by allowing us to choose time unit we want to display execution time in: seconds, mili-seconds or nano-seconds! In other words, we want to implement performance_log decorator in a way that will allow us to pass it as a parameter one of: "s", "ms", "ns", like this:

@performance_log(time_units="ms")
def some_function():
    pass

How can we achieve this? Let's start from the original performance_log decorator:

import time

def performance_log(func):

  def decorator(*args, **kwargs):

    start = time.perf_counter()
    result = func(*args, **kwargs)
    end = time.perf_counter()
    print(f"Execution time is: {end-start}")
    return result

  return decorator

We know that when Python interpreter sees the following lines:

@performance_log
def some_function():
    pass

it does the following:

some_function = performance_log(some_function)

And it is expected that the return type of the line performance_log(some_function) will be a function that receives another function as parameter and returns a decorated function.

So now, similarly, we need to update our performance_log such that

performance_log(some_function)(time_units="ms")

will return a function (the same one as before). Let's try implementing one:

# this should return a function that receives another function 
# and returns a decorated original function
def performance_log(time_units):

  # lets create it
  def wrapper(func):
    pass

  # now lets return it
  return wrapper

So far so good. We only need to add an implementation to this returned wrapper function. According to the requirement we mentioned above, this wrapper function should receive a function as parameter, and return its decorated function. Ok, let's do it:

# this should return a function that receives another function 
# and returns a decorated original function
def performance_log(time_units):

  # lets create one
  def wrapper(func):

    def decorator(*args, **kwargs):

      # code that performs some stuff and calls func received as param
      pass

    # returning our decorator  
    return decorator

  # now lets return one
  return wrapper

We are almost there! It's only left to complete the implementation of the actual decorator (it will be very similar to the original implementation, but will also take into account provided time_units parameter)

import time

def performance_log(time_units='s'):

  def wrapper(func):

    def decorator(*args, **kwargs):

      time_func = time.perf_counter if time_units != 'ns' else time.perf_counter_ns

      start = time_func()
      result = func(*args, **kwargs)
      end = time_func()
      exec_time = end - start
      if time_units == 'ms':
        exec_time *= 1000
      print(f"Execution time is: {exec_time}{time_units}")
      return result

    return decorator

  return wrapper

As you probably noticed, implementing a decorator that is able to receive parameters required us to add another "layer" of function definition, but as long as you understand the process behind the scenes, you won't experience any difficulties implementing stuff like this.

4. Implementing decorators inside a class

Until now we used to implement decorators as separate functions, but in fact it is possible and even recommended to implement them inside classes, especially if a decorator you implement is related to the code logic of the class.

Let's look at the example. We have the following Bank class:

import datetime
class Bank:

  def __init__(self, bank_name):
    self.name = bank_name

  '''
    Perform validation that the operation is being performed
    during working hours only: Sun - Thu, 09:00 - 17:00
  '''
  def working_hours_only(callable):

    def wrapped_callable(*args, **kwargs):
        pass

    return wrapped_callable

  @working_hours_only
  def withdraw(self, amount):
    print("Called withdraw", amount)
    return amount

  @working_hours_only
  def deposit(self, amount):
    print("Called deposit")

  def feedback(self, fedback_text):
    print("Called feedback")

As you can see, it implements 3 main methods: withdraw(), deposit() and feedback(). The first two are critical, hence we want to make sure they are not called outside working hours of the bank. The feedback() method is not critical, hence it can be called anytime. We want to implement a decorator working_hours_only that will decorate critical methods that should be called only during working hours of the bank. Since this decorator is tightly coupled with our Bank class logic, and is actually an inseparable part of it, it makes sense to implement this decorator inside a class.

One ❗️ important ❗️technical detail you should remember when defining decorator inside a class is that decorator signature inside a class should not contain self parameter, since it's not passed to the decorator by Python interpreter. And if you think about it, it totally makes sense. Why? Because defining class methods, including "replacing" original methods with decorated ones is happening during class definition stage, when there is no class instance present in the picture. You can think about decorators as of static methods defined inside a class. That's why we don't expect and don't need self parameter passed there.

After we discussed implementing decorators inside a class, you are ready to implement this working_hours_only decorator.

Exercise 2: Implement working_hours_only decorator inside Bank class

Implement decorator working_hours_only that validates that method is called during working hours only (Sun - Thu, 09:00 - 17:00). For example, calling feedback() method on Saturday should reach feedback() method and "Called feedback" should be printed out. Calling withdraw() method on Saturday should not reach actual withdraw() method and should raise an exception.

Here you can check out my solution for this exercise.

Exercise 3: Change your implementation of working_hours_only decorator to receive bank working days and hours as parameters.

I'm not publishing a solution for this exercise, but I encourage you to solve it and publish your own solutions in the comments. I'll be happy to mention the best solution here including credits for the author 😄

I hope you enjoyed this article. The next one will be published soon, stay tuned and subscribe to my channel to get notified as soon as it comes out!

All the code in Google Colab can be found here