Search code examples
pythondecoratorpython-decorators

How to decorate inherited method - Python


I have to decorate a inherited method, but it decorates all inherited methods. Basically I have to create a decorator that will decorate just one method from the class.

The test looks like this

@my_decorator
    class TestClass(Subclass):
        pass

t = TestClass()
t.say_hi

Let's say my SubClass looks like this

class SubClass():
    def __init__(self):
        pass

    def say_hi():
        print("Hi")

    def say_wow():
        print("wow")

Now I have to make my_decorator, that has to decorate inherited function say_hi() to print("*****") before it prints "Hi"

I tried doing it like this, but than the decorator applies to all methods from SubClass

def my_decorator(cls)
    def say_hi():
        print("*********")
        cls.say_hi()

    return say_hi()

Naturally It applies to every function of the subclass, but how do I make it to apply to only a say_hi() function? -It also returns an TypeError "NoneType" object is not callable


Solution

  • First let us fix SubClass, because instance methods require an explicit instance parameter at definition time:

    class SubClass():
        def __init__(self):
            pass
    
        def say_hi(self):
            print("Hi")
    
        def say_wow(self):
            print("wow")
    

    Now you want the decorator to replace the say_hi method with a method that prints '****' before calling the original method. Le us write a decorator that just does that(*):

    def my_decorator(cls):
        orig = cls.say_hi          # save the original method
        def say_hi(self):          # define a new one
            print('****')
            return orig(self)      # ... calling the original method
        cls.say_hi = say_hi        # replace the method in the class
        return cls
    

    You can then use:

    @my_decorator
    class TestClass(SubClass):
        pass
    t = TestClass()
    t.say_hi()
    

    and get as expected:

    ****
    Hi
    

    (*) this is a very simple decorator that can only replace a say_hi(self) method: neither a different name, nor additional parameters, but decorators can be much smarter...