Using decorators with class

Question

I am trying to add a decorator to required class methods and I have come up with the following code for it. I need this to work with all the similar classes.

import allure

def class_method_dec(cls):
    """
    Decorator function to decorate required class methods.
    """
    if cls.method1:
        cls.method1= allure.step(cls.method1)

    if cls.method2:
        cls.method2= allure.step(cls.method2)

    if cls.method3:
        cls.method3= allure.step(cls.method3)

    return cls


@class_method_dec
class TestClass:

    def __init__(self, a, b):
        self.a = a
        self.b = b

    def method1(self):
        """
        method docstring
        """
        pass

    def method2(self):
        """
        method docstring
        """
        pass

    def method3(self):
        """
        method docstring
        """
        pass

Is this the right way to do it? I am looking for the best way to do this.

Also, I understand that we can use functools.wraps to preserve the docstring when decorating functions. Is there a need of something like it when we are decorating classes?

Answer 1

From Satwik Kansal’s brilliant Metaprogramming in Python IBM tutorial , I discovered this gem:

Satwik first defined a decorator:

from functools import wraps
import random
import time

def wait_random(min_wait=1, max_wait=30):
    def inner_function(func):
        @wraps(func)
        def wrapper(args, **kwargs):
            time.sleep(random.randint(min_wait, max_wait))
            return func(args, **kwargs)

        return wrapper

    return inner_function

And then he created a class wrapper that will apply this decorator to a class:

def classwrapper(cls):
    for name, val in vars(cls).items():
        #callable return True if the argument is callable
        #i.e. implements the __call
        if callable(val):
            #instead of val, wrap it with our decorator.
            setattr(cls, name, wait_random()(val))
    return cls

Application:

# decorate a function

@wait_random(10, 15)
def function_to_scrape():
    #some scraping stuff

# decorate a class

@classwrapper
class Scraper:
    # some scraping stuff

To make use of it in your case, substitute wait_random decorator with your own. Turn your function to a decorator. E.g

from functools import wraps
import allure

def apply_allure():
    def inner_function(func):
        @wraps(func)
        def wrapper(args, **kwargs):
            func = allure.step(func)
            return func(args, **kwargs)

        return wrapper

    return inner_function

In the classwrapper replace wait_random with apply_allure:

Do read the tutorial for more information and explanations