In Python, decorators can either be functions or classes. In both cases, decorating adds functionality to existing functions. When we decorate a function with a class, that function becomes an instance of the class. When we define methods in the decorating class, we can add functionality to the function. This can all be achieved without modifying the original function source code.
This tutorial will demonstrate how classes can be used to decorate the functions we write in our Python code. We’ll go over two use-cases: Decorating a function with a class that accepts no arguments and decorating a function with a class that can accept arguments. If no arguments are passed, our class will fall back to a default value.
What Is a Python Class Decorator?
I have simplified the examples presented to make it easier to understand. Similar outcomes could be achieved using alternative, simpler approaches, but my goal in this tutorial is to demonstrate how we can use classes to decorate functions and extend their functionality.
This tutorial was also motivated to showcase how to use decorated classes that can accept arguments themselves.
How to Create a Python Class Decorator
The multiply_together function is designed to take two integer values, multiply them together and return their output. Let's consider a scenario in which we would like to add some extra functionality to the function without changing its original source code. Perhaps we would like to square the returned value. We can achieve this using a class decorator.
To decorate a function with a class, we must use the @ syntax followed by our class name above the function definition. Following convention, we will use camel case for our class name. In the class definition, we define two methods — the init constructor and the magic (or dunder) call method.
When we decorate a function with a class, the function is automatically passed as the first argument to the init constructor. We set this function as an attribute in our object. If we print multiply_together now, we can see it is an instance of the power class.
By defining the __call__() method, we can call multiply_together as you could with the original function. Here, we can see that we multiply two-by-two and square the answer.
The call method requires two arguments, which are specified because our original multiply together function required two arguments. We call the multiply together function with these two arguments in this method. This function has been set as self._arg in the object.attribute syntax below. We call this function with two values passed and save the returned value to the variable retval. Finally, we square retval and return the value.
The power class extends the functionality of the original multiply_together function. The source code for this example is shown below:
class Power(object):
def __init__(self, arg):
self._arg = arg
def __call__(self, a, b):
retval = self._arg(a, b)
return retval ** 2
@Power
def multiply_together(a, b):
return a * b
print(multiply_together)
print(multiply_together(2, 2))
Extending Multiply_Together’s Functionality By Adding Memory
To extend the example presented in the previous section, we can give our power object some memory of the squared values it returned. We can set an empty list to the memory attribute of our object, and append this list every time we call the decorated function. Finally, we can define a method, named “memory” in the code below to return the values stored in the list held by the memory attribute.
This way, we have extended the functionality of our multiply together function further.
The underscore _memory attribute now stores the list of squared values passed to multiply together. We simply call the memory method on the power instance to retrieve the result.
Python Class Decorators That Can Accept Arguments
To increase the functionality of the example even further, it would be better to have our class decorator accept arguments. In this way, we could choose which value we would like to use as the exponent with our power class. If no argument is passed to the class decorator, a default exponent value will be set.
How to Pass Arguments to the Python Class Decorator
When we pass an argument to the class decorator, that argument and not the function is passed as the argument to the init constructor. In the example presented, we pass the integer value three as an argument to the power class constructor. This value is saved as an attribute, underscore arg (_arg) in the object. The function is then passed as the only argument when we define the call method.
Therefore, if the length of the arguments passed to the call method is equal to one, as would be the case if we pass a decorator argument to the class, this first argument to the call method will be set as the function. We can then define an inner function inside the call method that takes two arguments, “a” and “b.” The call method returns the wrapper function if the length of the arguments passed to call is one. We can call this function with two values passed and finally multiply it by the integer (stored under the attribute _arg) that was originally passed to the class as an argument.
We can then use asterisks followed by the parameter name, here, param_arg, to add flexibility to our call method. This means the parameter can accept a variable number of arguments which are stored in a tuple and allows length checking.
To aid this description, I have included the corresponding example with the accompanying document string below.
How to Pass No Arguments to the Python Class Decorator
The alternative scenario is when no argument is passed to the class decorator. In this particular case, the function is passed as the first argument to the init constructor.
When we call our decorated function, we pass two integer arguments. Thus, the first conditional in the call method fails and execution proceeds to the else statement. Here, a default value of two is set as the exponent, and the function, stored in the underscore arg attribute of the object, passes the two arguments. The return value is stored in the variable retval. Finally, retval is multiplied by the default exponent and returned.
The source code for both examples shown can be found below.
Now we’ve satisfied both conditions — a class decorator that can either accept arguments or not. If no argument is passed to the class decorator, a default can be set.
# Full example version with accompanying doc string
class Power(object):
def __init__(self, arg):
self._arg = arg
def __call__(self, *param_arg):
"""If there are decorator arguments, __call__() is only called once
as part of the decoration process. You can only give it a single argument,
which is the function object
If there are no decorator arguments, the function
to be decorated is passed to the constructor.
"""
if len(param_arg) == 1:
def wrapper(a, b):
retval = param_arg[0](a, b)
return retval ** self._arg
return wrapper
else:
expo = 2
retval = self._arg(param_arg[0], param_arg[1])
return retval ** expo
# @Power(3)
# def multiply_together(a, b):
# return a * b
@Power
def multiply_together(a, b):
return a * b
print(multiply_together(2, 2))
Python Class Decorator Advantages
The output from the code above could be achieved through simpler means, however, this article focuses on how to use class decorators. Therefore, my goal has been to provide easy-to-follow examples.
Functions can be decorated with classes to extend their functionality. Further, classes that decorate functions can either accept arguments or fall back to a default if no argument is passed. Here, both use-cases are presented to improve the functionality of the original function.
