Simulate method_added Hook in Python

Recently, for whatever the reason, I've ended up reading some stuff about the amazing Ruby metaprogramming capabilities (I have no Ruby experience, I just toyed with it more than 15 years ago, and from then I've occasionally taken a fast peek to compare it with my known languages). Ruby hooks seem pretty amazing to me. I'd never seen anything similar in other languages and are one of those language features that after the initial amazement you wonder, how would I use them, and how could I simulate them in other languages?

I'm particularly interested in the method_added hook. So you can hook to your class a function that will be executed each time a method is added to that class (both if the method is defined as part of the class or if it's added later to the class (aka expando methods)). Amazing, but what's a practical use for that? Well, main thing that comes to my mind is AOP. Let's say we want to add to all methods in one class a call to log("method started") and log("method finished"). We would hook to method_added a function that receives the original function and creates a new function that wraps the call to the original function between those calls to the logger. And now the big question, how could we achieve something similar in Python? Well, creating wrapper functions... that's what decorators are all about. So first let's create a function that when invoked creates a new wrapper function that performs this extra logging.

# wraps a function in a new function that performs logging before and after invokation
def log_wrapper_creator(fn: Callable) -> Callable:
    @wraps(fn)
    def log_added(*args, **kwargs):
        print(f"{fn.__name__} started")
        res = fn(*args, **kwargs)
        print(f"{fn.__name__} finished")
        return res
    return log_added
	

An now let's create a decorator that when applied to a class, applies the above decorator to all methods in the class.

# decorator that applies to a class wrapping all its methods using the provided wrapper_creator function
def wrap_methods(wrapper_creator: Callable):
    def _wrap_methods(cls):
        members = inspect.getmembers(cls, predicate=inspect.isfunction)
        for name, fn in members:
            setattr(cls, name, wrapper_creator(fn))
        return cls
    return _wrap_methods 
	

That's nice, but it's only one part of what we get with Ruby method_added. If after defining our class then we decide to add a new method to it (we monkey-patch or expand the class with new methods) the method_added hook will also handle that. So how can we emulate that in Python? __setattr__ is a sort of "generic hook" that we can use in Python to react to the setting of an attribute in one object. If we want to react to setting attibutes in a class, we'll need __setattr__ defined in its metaclass. We can add extra logic to our existing decorator, so that after wrapping the existing methods it sets the metaclass of our class to a metaclass with _setattr__ powers. Important point, as I explain at the end of this post the metaclass of a class can not be changed (__class__ is readonly in instances of type), so we have to use an extra trick. We can define a new class that extends the existing class and has our custom metaclass as metaclass. Yeah, much talk, let's see in action what I mean:

# Metaclass
class WrapMeta(type):
    def __setattr__(cls, name, value):
        if inspect.isfunction(value):
            value = cls._wrapper_creator(value)
        return super().__setattr__(name, value)

def wrap_methods_dynamic(wrapper_creator: Callable):
    def _wrap_methods(cls):
        # wrap method that exist when the class is provided
        members = inspect.getmembers(cls, predicate=inspect.isfunction)
        for name, fn in members:
            setattr(cls, name, wrapper_creator(fn))
        # and use a metaclass for methods that will be added in the future
        # create a new class extending the main one and using as metaclass WrapMeta
        class A(cls, metaclass=WrapMeta):
            pass
        # this one fails cause it invokes the __setattr__ that tries to access the wrapper_creator that we are trying to set
            # A.wrapper_creator = wrapper_creator 

        # so we have to use:
            # object.__setattr__ does not work with classes ( an't apply this __setattr__ to WrapMeta object), we have to use type.__setattr__
            # object.__setattr__(A, "wrapper_creator", wrapper_creator)
        type.__setattr__(A, "_wrapper_creator", wrapper_creator)      
        A.__name__ = cls.__name__ 
        return A
    return _wrap_methods 

Notice how for the __setattr__ in the metaclass to be able to apply our wrapper_creator function, we have to make it accessible somewhere. We do so but setting it as a "private" attibute of the class, and that's a bit tricky. Setting that attribute will trigger our custom __setattr__ so to prevent that, we'll have to specify that we want to use a different __setattr__ logic. In a method of a class, we use super().__setattr__, but here we are in a function, so I tried to use object.__setattr__. That fails with a TypeError: can't apply this __setattr__ to WrapMeta object. What we have to use is type.__setattr__. It's discussed here.

type.__setattr__ is used for classes, basically instances of metaclasses. object.__setattr__ on the other hand, is used for instances of classes.

Ah, of course, let's see an example:

@wrap_methods_dynamic(log_wrapper_creator)
class User:
    def __init__(self, name):
        self.name = name

    def say_hi(self, sm: str):
        return f"{self.name} says hi to {sm}"

    def walk(self):
        return f"{self.name} is walking"


u1 = User("Iyan")
print(u1.say_hi("Francois"))
print(u1.walk())
User.do_work = lambda self, a, b: f"{self.name} is working on {a} and {b}"
print(u1.do_work("aaa", "bbb"))

# __init__ started
# __init__ finished
# say_hi started
# say_hi finished
# Iyan says hi to Francois
# walk started
# walk finished
# Iyan is walking
#  started
#  finished
# Iyan is working on aaa and bbb