Adding hooks to functions in subclassed methods

Viewed 677

Given the following simplified code:

from abc import ABC, abstractmethod

class Parent(ABC):
    def __init__(self,*args,**kwargs):
        self.parent_name = 'SuperClass'
    
    # global hook to run before each subclass run()
    def global_pre_run_hook(self):
        pass
    
    @abstractmethod
    def run(self, *args, **kwargs):
        raise NotImplementedError()

        
        
class Child(Parent):
    def __init__(self,*args,**kwargs):
        super().__init__(*args,**kwargs)
        self.name = 'ChildClass'
    
    def run(self):
        print(f'my parent name is {self.parent_name}')
        print(f'my name is {self.name}')
        
        return 22
        
obj = Child()
result = obj.run()

Is there a way to add functionality so that when the child class run() method is called directly, it first triggers a hook function from the parent class? Assume there is a parent class and a lot of classes that subclass it - would I need to manually add a call global hook() at the beginning of each run() definition for each class that subclasses Parent()? Is there a pythonic way to accomplish this?

2 Answers

There might be a way to do this with a proper decorator, but for the time being I think the cleanest solution you might come up with would be to create your own 'decorator' and manually apply it in the course of Parent.__init__(), which would make sure it takes effect so long as the child class invokes the parent __init__():

from abc import ABC, abstractmethod

def create_hook(func, hook):
    def wrapper(*args, **kwargs):
        hook()
        return func(*args, **kwargs)
    return wrapper

class Parent(ABC):
    def __init__(self, *args, **kwargs):
        self.parent_name = 'SuperClass'
        self.run = create_hook(self.run, self.global_pre_run_hook)

    # global hook to run before each subclass run()
    def global_pre_run_hook(self):
        print("Hooked")

    @abstractmethod
    def run(self, *args, **kwargs):
        raise NotImplementedError()

class Child(Parent):
    def __init__(self,*args,**kwargs):
        super().__init__(*args,**kwargs)
        self.name = 'ChildClass'

    def run(self):
        print(f'my parent name is {self.parent_name}')
        print(f'my name is {self.name}')
        return 22

obj = Child()
result = obj.run()
# this prints:
#   Hooked
#   my parent name is SuperClass
#   my name is ChildClass

The answer by Green Cloak Guy works, but cannot be pickled! To fix this, we need to move the hook creation into __new__. Also, it's a good idea to make use of functools.wraps in the hook creator.

import pickle
from abc import ABC, abstractmethod
from functools import wraps

def create_hook(func, hook):
    @wraps(func)
    def wrapper(*args, **kwargs):
        hook(*args, **kwargs)
        return func(*args, **kwargs)
    return wrapper

class Parent(ABC):
    def __new__(cls, *args, **kwargs):
        cls.run = create_hook(cls.run, cls.global_pre_run_hook)
        return super().__new__(cls, *args, **kwargs)

    # global hook to run before each subclass run()
    def global_pre_run_hook(self, *args, **kwargs):
        print("Hooked")

    @abstractmethod
    def run(self, *args, **kwargs):
        raise NotImplementedError()

class Child(Parent):
    def run(self):
        print(f"my parents are {self.__class__.__mro__}")
        print(f"my name is {self.__class__.__name__}")
        return 22

obj = Child()
result = obj.run()
pickle.dumps(obj)
Related