The Class Decorator in Python

1. Class Decorator in Python
2. Extend the Functionality of Code in Python
3. Arguments With Class Decorators in Python
4. Use *Args and **Kwargs as Arguments in Python
5. Decorator That Has a Return Statement
6. Get the Execution Time in Python
7. Use Class Decorator to Check Error Parameter in Python
8. Conclusion

In Python, we can extend the behavior of a function or a class without modifying it.

We can wrap functions inside other functions to add some functionality to the existing class or function with the help of decorators.

Class Decorator in Python

Decorators is a tool in Python that lets the programmer modify the behavior of a class or function.

We can visualize the decorators as a three-step process where:

1. We give some function as an input to the decorator.
2. The decorator works to add functionality.
3. The decorator returns the function with added usage.

For example, we have function A, and we want to add functionalities without modifying them permanently. We can use a decorator as a class by using the __call__ method.

Callable is any object that can implement the __call__ method. A decorator is a callable that can return a callable.

In layman language, if an object is similar to a function, the function decorator should also return an object similar to a function. Here’s an example using the __call___ method.

# Use of the __call__ method in Python
class MyDemoDecorator:
    def __init__(self, func):
        self.func = func

    def __call__(self):
        # Code before the function call
        self.func()
        # Code after the function call


# adding the decorator as class
@MyDemoDecorator
def func():
    print("Learning!")


func()

Output:

Learning

When we decorate a function using class, we make that function an instance of the decorating class.

The class we use for decorating a function can have arguments, but if we do not pass any argument, the class falls back to the default value.

Extend the Functionality of Code in Python

We have a function mul_nums() that multiplies two numbers and returns the product as the output. Now, we want this function to return the product and the cube of the product as well.

The part that calculates the product’s cube is an extra function that we will not add to the source code. Rather, we will use a class decorator to achieve this added functionality.

We decorate the function with class using @Cube in the code block below.

Example:

class Cube(object):
    def __init__(self, args):
        self.args = args

    def __call__(self, x, y):
        res = self._args(x, y)
        return res * res * res


@Cube
def mul_nums(x, y):
    return x * y


print(mul_nums)
print(mul_nums(4, 3))

Output:

1728

The init constructor inside the class automatically receives the function as the first argument. The function is set as an attribute inside the object.

Therefore, we can see the function mul_nums() as an instance of the Cube class when we print mul_nums.

Inside the __call__() method, we call the mul_nums function where multiplication and cubing of result occur. The value is returned after finding its cube.

There is one more function that we can add to this code. Suppose we give some memory of the cubed value to our cube object.

For this, we use an empty list and set it to the attribute corresponding to the object’s memory. Every time we call the decorated function, we also append it to this list.

Lastly, we define the method mem, which returns the stored values from the list.

Example:

class Cube(object):
    def __init__(self, args):
        self._args = args
        self._mem = []

    def __call__(self, x, y):
        res = self._args(x, y)
        self._mem.append(res * res * res)
        return res * res * res

    def mem(self):
        return self._mem


@Cube
def mul_nums(x, y):
    return x * y


print(mul_nums)
print(mul_nums(4, 3))
print(mul_nums(2, 3))
print(mul_nums(5, 2))
print(mul_nums.mem())

Output:

1728

Arguments With Class Decorators in Python

A class decorator has two types. One accepts arguments, and the other does not.

Both types work fine, but the class decorator that can take an argument is more flexible and efficient.

Let us see both cases one by one. This time we will look at a scenario where we define a function add_num() that adds two numbers.

Then, we will use the concept of class decorator and the functionality of add_num() to get the power of the result. In the below example, the class decorator takes one argument.

class Power(object):
    def __init__(self, args):
        self._args = args

    def __call__(self, *param_arg):
        if len(param_arg) == 1:

            def wrap(x, y):
                res = param_arg[0](x, y)
                return res ** self._args

            return wrap
        else:
            exponent = 2
            res = self._args(param_arg[0], param_arg[1])
            return res ** exponent


@Power(2)
def add_num(x, y):
    return x + y


print(add_num(4, 3))

Output:

49

Here, the init function does not get the function as an argument. Rather, the argument that we pass to the class decorator goes to the init constructor.

The value 2 that we pass here as an argument is saved as an attribute. Later, when we define the __call__ method, the function is the only argument passed there.

Note that if the length of the arguments we pass to the __call__ method is 1, the method returns the wrap function. The use of asterisks * with param_arg is to allow a variable number of arguments.

Now let us look at the alternative case where we do not pass any argument to the class decorator.

class Power(object):
    def __init__(self, args):
        self._args = args

    def __call__(self, *param_arg):
        if len(param_arg) == 1:

            def wrap(x, y):
                res = param_arg[0](x, y)
                return res ** self._args

            return wrap
        else:
            exponent = 2
            res = self._args(param_arg[0], param_arg[1])
            return res ** exponent


@Power
def add_num(x, y):
    return x + y


print(add_num(4, 3))

Output:

49

Since no argument is passed to the class decorator, the init constructor gets a function as the first argument. Calling the decorated functions fails the first condition, and therefore, the else block executes.

Inside the else block, a default value is set. Using this default value, we get the resultant value.

Use *Args and **Kwargs as Arguments in Python

In the example above, the __call__ function takes one argument. Another way of using the class decorator is to pass the arguments *args and **kwargs in this function.

Example:

class MyDemoDecorator:
    def __init__(self, func):
        self.func = func

    def __call__(self, *args, **kwargs):
        # code before the function call
        self.func(*args, **kwargs)

    # code after the function call


# adding class decorator to the function
@MyDemoDecorator
def func(name, msg="Hey there"):
    print("{}, {}".format(msg, name))


func("we are learning decorators", "hey there")

Output:

hey there, we are learning decorators

Decorator That Has a Return Statement

Let us work with functions that do return some value.

In such cases, we use the return statement.

Example:

# decorator having a return statement
class DemoDecorator:
    def __init__(self, func):
        self.func = func

    def __call__(self, *args, **kwargs):
        # code before function call
        res = self.func(*args, **kwargs)
        # code after the function call
        return res


# adding class decorator to the function
@DemoDecorator
def cube(n):
    print("The given number is:", n)
    return n * n * n


print("Cube of the given number is:", cube(11))

Output:

The given number is: 11
Cube of the given number is: 1331

Get the Execution Time in Python

We can use the class decorator to print a program’s time for execution. Use the __call__() function with the time module.

Example:

# using class decorator to get the execution time of a program
# import the time module
from time import time


class Execution_Time:
    def __init__(self, func):
        self.funct = func

    def __call__(self, *args, **kwargs):
        start_time = time()
        res = self.funct(*args, **kwargs)
        stop_time = time()
        print(
            "The execution of this program took {} seconds".format(
                stop_time - start_time
            )
        )
        return res


# adding decorator to a function
@Execution_Time
def demo_function(delay):
    from time import sleep

    # delaying the time
    sleep(delay)


demo_function(3)

Output:

The execution of this program took 3.004281759262085 seconds

Use Class Decorator to Check Error Parameter in Python

One of the class ’ decorator ’ uses is checking the parameters of a function before executing. It prevents the function from overloading, and only logical, and most necessary statements are stored.

Example:

# use class decorator to check error parameter
class CheckError:
    def __init__(self, func):
        self.func = func

    def __call__(self, *params):
        if any([isinstance(i, str) for i in params]):
            raise TypeError("Parameter is a string and it ain't possible!!")
        else:
            return self.func(*params)


@CheckError
def add(*numbers):
    return sum(numbers)


#  calling function with integers
print(add(3, 5, 2))
#  calling function with a string in between
print(add(3, "5", 2))

Output:

10
TypeError: Parameter is a string and it ain't possible!!

Conclusion

We discussed the concept and use of Python class decorators. We also discussed how a class decorator could return statements, get the execution and check error parameters.