Search code examples
pythonpython-3.xdecoratorpython-decorators

How to wrap a decorator around another classes method?


I have created a decorator which I am using to manage logging. I want logging to occur before and after the decorated function runs. The function works fine when interacting with very basic functions, however, when interacting with methods that are a part of other classes, things break. I suspect the issue is a result of there being 2 self arguments. Do you have any idea how to resolve it?

Simplified Decorator Class

class Logger:

    def __init__(self, logging_type:str = 'debug'):
        self.logging_type = logging_type

    def __call__(self, decorated_function:callable):
        self.func = decorated_function
        return getattr(self, self.logging_type)

    def debug(self, *args, **kwargs):
        print("starting function")
        output = self.func(*args, **kwargs)
        print("Completing Function")
        return output

We see that the decorator works on basic functions:

@Logger(logging_type="debug")
def simple_function(x):
    return x**2

In [2]: simple_function(3)
starting function
Completing Function
Out[2]: 9

However, fails when work with other classes:

class BigClass:

    def __init__(self, stuff = 10):
        self.stuff = stuff

    @Logger(logging_type="debug")
    def cool_function(self, input1: int):
        return self.stuff + input1


In [16]: test = BigClass()
    ...: test.cool_function(3)
starting function

It then hits a type error on the output line:

TypeError: cool_function() missing 1 required positional argument: 'input1'

Ideas?


Solution

  • By all means read juanpa.arrivillaga's informative answer. But here is a simpler approach. In writing a class decorator of this type, __call__ should return an ordinary function instead of a member function, like this:

    class Logger:
        def __init__(self, logging_type:str = 'debug'):
            self.logging_function = getattr(self, logging_type)
    
        def __call__(self, decorated_function: callable):
            def f(*args, **kwargs):
                return self.logging_function(decorated_function, *args, **kwargs)
            return f
    
        def debug(self, decorated_function, *args, **kwargs):
            print("starting function")
            output = decorated_function(*args, **kwargs)
            print("Completing Function")
            return output
    
    @Logger(logging_type="debug")
    def simple_function(x):
        return x**2
    
    class BigClass:
        def __init__(self, stuff = 10):
            self.stuff = stuff
    
        @Logger(logging_type="debug")
        def cool_function(self, input1: int):
            return self.stuff + input1
    
    print(simple_function(12))
    
    test = BigClass()
    print(test.cool_function(3))
    

    OUTPUT:

    starting function
    Completing Function
    144
    starting function
    Completing Function
    13