How to check for class attribute initialization with a decorator?-CodePudding

Consider the following code.

def check(data):
    def decorator(function):
        def wrapper(*args, **kwargs):
            if data not in ["toto", "tata"]:
                raise ValueError("Wrong argument")
            result = function(*args, **kwargs)
            return result
        return wrapper
    return decorator


class Foo:
    
    @check
    def __init__(data):
        self.data = data
         
         
Foo(data="t")

I would like to decorate the __init__ function of a class, but I get an exception :

Traceback (most recent call last):
  File "<string>", line 19, in <module>
TypeError: decorator() got an unexpected keyword argument 'data'

Any help would be appreciated :)

CodePudding user response：

You're mixing up which function receives what. check is the decorator which receives the function you're decorating, not data. wrapper is the function which will receive data. You don't need the third function in there. So:

def check(function):
    def wrapper(self, data):
        if data not in ["toto", "tata"]:
            raise ValueError("Wrong argument")

        return function(self, data)

    return wrapper