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

This is a 2-Part 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

---

This article is Part 2 in a series about Python Decorators, and it assumes you have basic knowledge of 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, specifically:

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

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

Let’s start with an example that shows why we need decorators.

Recap

Let’s begin with a short recap of what a decorator is and what we learned in the 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.

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

greeting_decorator(sum_of_digits)

When we call:

sum_of_digits()

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

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

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.

If we try running:

sum_of_digits(235)

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

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

greeting_decorator(sum_of_digits)(235)

This means that we need to change greeting_func returned from greeting_decorator to accept arguments and return a value, as follows:

Now, when calling:

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

we will get the expected output:

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

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

Now, we can use our greeting_decorator with any function!

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

Now it’s time for an exercise. I encourage you to try solving it 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 the amount of time (in seconds) that it takes for a function to complete execution. This 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.

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 yourself 😅

---

Spoiler alert⚠️: The next section is based on the performance_log decorator.

Passing Parameters to Decorators

After we implemented the performance_log decorator that logs function execution time in seconds, we want to add more flexibility by allowing us to choose the time unit we want to display execution time in: seconds, milliseconds, or nanoseconds. In other words, we want to implement the performance_log decorator in a way that will allow us to pass it a parameter, such as one of: “s”, “ms”, or “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:

We know that when the 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 a parameter and returns a decorated function.

Now, similarly, we need to update our performance_log so that it will return a function (the same one as before). Let’s try implementing one:

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

We are almost there! We just need to complete the implementation of the actual decorator (it will be very similar to the original implementation, but will also take into account the provided time_units parameter).

As you probably noticed, implementing a decorator that is able to receive parameters required us to add another “layer” of function definition. However, as long as you understand the process behind the scenes, you won’t experience any difficulties implementing such features.

Implementing Decorators Inside a Class

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

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

As you can see, it implements three main methods: withdraw(), deposit(), and feedback(). The first two are critical, hence we want to make sure they are not called outside the working hours of the bank. The feedback() method is not critical, so 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 the working hours of the bank. Since this decorator is tightly coupled with our Bank class logic, and is an inseparable part of it, it makes sense to implement this decorator inside the class.

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

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

Exercise 2: Implement the working_hours_only Decorator Inside the Bank Class

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

Here you can check out my solution for this exercise.

Exercise 3: Modify Your Implementation of the 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, so 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.