Search code examples
pythonclassoopdecoratorpython-decorators

Apply decorator to class method which accesses class attribute


Is it possible to write a decorator that acts upon a class's method and uses the class's attributes? For example, I would like to add a decorator to functions that will return an error if one of the class's attributes (which is set when the user calls the function) is False.

For example, my attempt (broken code since is_active can't access MyClass's methods):

def is_active(active):
    if active == False:
        raise Exception("ERROR: Class is inactive")


class MyClass():

    def __init__(self, active):
        self.active = active

    @is_active
    def foo(self, variable):
        print("foo")
        return variable

    @is_active
    def bar(self, variable):
        print("bar")
        return variable

where the expected behaviour is:

cls = MyClass(active=True)
cls.foo(42)
---> function prints "foo" and returns 42

cls = MyClass(active=False)
cls.foo(42)
---> function raises an exception as the active flag is False

The above is a dummy example and the actual use case is more complex, but hopefully this shows the problem I'm facing.

If the above is possible, my extra question is: is it possible to "hide"/delete the methods from the instantiated class based on this flag. For example, if the user instantiates the class with a active=False then when they're using iPython and press <tab>, they can only see the methods which are permitted to be used?

Thank you.


Solution

  • Decorators can be confusing. Note a function is passed as a parameter and the decorator expects that a function (or callable object) is returned. So you just need to return a different function. You have everything else you need since self is passed as the first argument to a class method. You just need to add a new function in your decorator that does what you want.

    def is_active_method(func):
        def new_func(*args, **kwargs):
            self_arg = args[0] # First argument is the self
            if not self_arg.active:
                raise Exception("ERROR: Class is inactive")
            return func(*args, **kwargs)
        return new_func
    
    
    class MyClass():
    
        def __init__(self, active):
            self.active = active
    
        @is_active_method
        def foo(self, variable):
            print("foo")
            return variable
    
        @is_active_method
        def bar(self, variable):
            print("bar")
            return variable
    
    m = MyClass(True) # Prints foo from the method
    m.foo(2)
    
    m = MyClass(False) # Outputs the exception
    m.foo(2)