pythonpython-decoratorsselfpython-class

Python decorators with class methods


I want a decorator that I can use as both @decorator and decorator() with class methods, like this :

def decorator_with_args(name):
    print(f'Hello {name} !')
    def decorator(func):
        def wrapper(self, *args, **kwargs):
            print(f'Hello {self.title} {name} again !!')
            print(f"Calling {func.__name__} with instance {self}...")
            result = func(self, *args, **kwargs)
            print(f"{func.__name__} finished. Result: {result}")
            return result
        return wrapper
    return decorator

class Calculator:
    def __init__(self):
        self.title = 'Mr'

    @decorator_with_args('World')
    def add(self, a, b):
        return a + b

    def add2(self, a, b):
        return a + b

    def do_add(self, a, b):
        return decorator_with_args(f'World{a}')(self.add2)(self, a, b)

# Usage 1
calc = Calculator()
result = calc.add(3, 5)

# Usage 2
calc = Calculator()
result = calc.do_add(3, 5)

The reason why I want to use the decorator function as the above two representations is because:

  1. I want to use @decorator_with_args when the argument is known such as @decorator_with_args('World')
  2. I want to use decorator_with_args() when the argument is not known and I want the argument to be dynamic such as decorator_with_args(f'World{a}')

While @decorator_with_args works as expected [Usage 1], I get an error with decorator_with_args() [Usage 2]. I tried few things, but none of them worked for the latter.

If I try to pass (self, a, b), for example decorator_with_args(f'World{a}')(self.add2)(self, a, b), I get TypeError: add2() takes 3 positional arguments but 4 were given.

On the other hand, if I try to pass (a, b) without self, for example decorator_with_args(f'World{a}')(self.add2)(a, b), I get AttributeError: 'int' object has no attribute 'title'.

I appreciate there may exist other similar questions, which I tried to search but could not get them to work for my use case.


Solution

  • In order to replicate the standard decorating of instance method, you need to decorate the underlying function that is bound to that method, not the method itself - that is because self.add2 implicitly adds self as first argument to that function, so it gets doubled.

        def do_add(self, a, b):
            return decorator_with_args(f'World{a}')(Calculator.add2)(self, a, b)
    

    or

        def do_add(self, a, b):
            return decorator_with_args(f'World{a}')(self.add2.__func__)(self, a, b)