Search code examples
pythondecoratorpython-decoratorsinstance-methods

in one class, applying decorator on a class method using another class method


I have a class Alpha that has method __A and method __B. they both called within the class __init__ scope or other Alpha method, thus the use of double_under __, and they both functions on Alpha attributes.

how can I make another wrapper method __C to work as a decorator to extend __A functionality and apply __B before and after __A execution.

the code:

class Alpha:
    def __init__(self, a, b) -> None:
        self.a = a
        self.b = b
        self.run_me()
        # partial(self.__C, self=self)

    #@staticmethod
    def __C(self, func):
        def wrapper(*args, **kwargs):
            # self: Alpha
            self.__B()
            func(*args, **kwargs)
            self.__B()
        return wrapper

    def __B(self):
        self.b +=5
        print('Method B ran')

    @__C
    def __A(self):
        self.a += 1
        print('Method A ran')

    def run_me(self):
        self.__A()

alpha = Alpha(1,5)

the comments are some ideas that I've tried, but I think I doing something wrong.

how can it be applied? what is wrong with my implementation?

note: I know I can manually rewrite __A, but in this example __A stands for a lot of methods that needs to be extended with __B functionality. thanks in advance.


Solution

  • To make __C work as a decorator for a method, it should be defined as a static method with the wrapper function inside taking self as a positional argument so that it can call __B as a bound method of the self instance:

    @staticmethod
    def __C(func):
        def wrapper(self, *args, **kwargs):
            self.__B()
            func(self, *args, **kwargs)
            self.__B()
        return wrapper
    

    Demo here