How to Decorate Inherited Methods Using a Custom Metaclass in Python?

I’m using a custom metaclass, to decorate methods of a class. However, it currently only decorates methods defined directly in the class and not those inherited from the parent class.

My classes are structured as follows:

class Abstract1(ABC):

class A(Abstract1):
   def fun1(self):
   def fun2(self):

class DecoMeta(ABCMeta):

class B(A, metaclass=DecoMeta):
   def hello(self):

My DecoMeta class looks like following:

class DecoMeta(ABCMeta):
    """ Decorator class
    """
    def __new__(mcs, name: str, bases: tuple, attrs: dict):
        for attr_name, attr_value, in attrs.items():
            if isinstance(attr_value, types.FunctionType):
                attrs[attr_name] = mcs.deco(attr_value)

        return super(DecoMeta, mcs).__new__(mcs, name, bases, attrs)

    @classmethod
    def deco(mcs, func):
        """ Decorate func
        """
        async def wrapper(*args, **kwargs):
            result = await func(*args, **kwargs)
            # print(args[0].__dict__.keys())
            args[0].some_function(). # i.e. self.some_function()
            return result
        return wrapper

When I print args[0].dict.keys() it only prints out the functions defined in class B. So that means, decoMeta only decorates those functions(in this case hello()). I want to decorate all functions defined in the parent class i.e. class A(like fun1() and fun2()) using this DecoMeta. Is there a way to do it?

Solution

You have to register A to the metaclass as well.

One way you can achieve that is by passing the metaclass to the parent class as well:

class Abstract1(ABC):
    pass

class DecoMeta(ABCMeta):
    pass

class A(Abstract1, metaclass=DecoMeta):
    def fun1(self):
        pass

    def fun2(self):
        pass

class B(A, metaclass=DecoMeta):
    def hello(self):
        pass

With the modified DecoMeta, just to display the attr_name:


class DecoMeta(ABCMeta):
    """ Decorator class
    """
    def __new__(mcs, name: str, bases: tuple, attrs: dict):
        for attr_name, attr_value, in attrs.items():
            if isinstance(attr_value, types.FunctionType):
                print(attr_name)
                # attrs[attr_name] = mcs.deco(attr_value)

        return super(DecoMeta, mcs).__new__(mcs, name, bases, attrs)

B()

outputs

fun1
fun2
hello

Otherwise, if you don't want to edit A at all, you can do something fancier like that:

class DecoMeta(ABCMeta):
    """ Decorator class
    """
    def __new__(mcs, name: str, bases: tuple, attrs: dict):
        all_attrs = attrs | {
            k: v for k, v in base.__dict__.items()
            for base in bases
        }

        for attr_name, attr_value, in all_attrs.items():
            if isinstance(attr_value, types.FunctionType):
                print(attr_name)
                # attrs[attr_name] = mcs.deco(attr_value)

        return super(DecoMeta, mcs).__new__(mcs, name, bases, attrs)

Which gives the same output.