从基类包装派生类方法

Question

Imagine I have a base and a derived class like so:

class A:
    def foo(self):
        pass

class B(A):
    def foo(self):
        pass

I wish to wrap calls foo calls made by instances of B . Modifying any part of B is not allowed (I don't own B).

What I have as of now:

class A:
    def __init__(self):
        fnames = ["foo"]
        for fname in fnames:

            def capture(f):
                def wrapper(*args, **kwargs):
                    print("wrapped")
                    return f(*args, **kwargs)

                return wrapper

            meth = capture(getattr(self, fname))
            bound = meth.__get__(self)
            setattr(self, fname, bound)


class B(A):
    def foo(self):
        print("b")


o = B()
o.foo()

print(o.foo)

This works as expected, however I am concerned about this being inefficient memory-wise.

o.foo is a <bound method A.__init__.<locals>.capture.<locals>.wrapper of <__main__.B object at 0x10523ffd0>> . It would seem that I would have to pay the cost of 2 closures for each instance I create.

Is there a better way to do this? Perhaps a metaclass based approach?

Answer 1

Unless you are planning to have several thousands instances of these live at the same time, the resource usage of doing so should not concern you - it is rather small compared with other resources a running Python app will use.

Just for comparison: asyncio coroutines and tasks are the kind of object which one can create thousands of in a process, and it is just "ok", will have a similar overhead.

But since you have control of the base class for B there are several ways to do it, without resorting to "monkey patching" - which would be modifying B in place after it was created. That is often the only alternative when one has to modify a class for which they don't control the code.

Wrapping the method either automatically, when it is retrieved from a B instance, in a lazy way, can spare even this - and sure can be more elegant than wrapping at the base class __init__ :

If you know beforehand the methods you have to wrap, and is sure they are implemented in subclasses of classes which you control, this can be made by crafting a specialized __getattribute__ method: this way, the method is wrapped only when it is about to get used.

from functools import wraps, partial

def _capture(f):  # <- there is no need for this to be inside __getattribute__
                  # unless the wrapper is to call `super()`
    @wraps(f)
    def wrapper(self, *args, **kwargs):
        print("wrapped")
        return f(*args, **kwargs) 
        # ^ "f" is already bound when we retrieve it via super().__getattribute__
        # so, take care not to pass "self" twice. (the snippet in the question
        # body seems to do that)
    return wrapper

class A:
    
    def __getattribute__(self, name):
        fnames = {"foo", }
        attr = super().__getattribute__(name)
        if name in fnames:
            # ^ maybe add aditional checks, like if attr is a method,
            # and if its origin is indeed in a
            # class we want to change the behavior
            attr = partial(_capture(attr), self)  
            # ^ partial with self as the first parameter
            # has the same effect as calling __get__ passing
            # the instance to bind the method
            
        return attr
            
class B(A):
    def foo(self):
        pass

---

As for wrapping foo when B is created, that could give use even less resources - and while it could be done in a metaclass, from Python 3.6 on, the __init_subclass__ special method can handle it, with no need for a custom metaclass.

However, this approach can be tricky if the code might further subclass B in a class C(B): which will again override foo : the wrapper could be called multiple times if the methods use super() calls to foo in the base classes. Avoiding the code in the wrapper to run more than once would require some complicated state handling (but it can be done with no surprises).

from functools import wraps

def _capture(f):
    @wraps(f)
    def wrapper(self, *args, **kwargs):
        print("wrapped")
        return f(self, *args, **kwargs) 
        # ^ "f" is retrieved from the class in __init_subclass__, before being 
        # bound, so "self" is forwarded explicitly
    return wrapper

class A:

    def __init_subclass__(cls, *args, **kw):
        super().__init_subclass__(*args, **kw)
        fnames = {"foo",}
        for name in fnames:
            if name not in cls.__dict__:
                continue
            setattr(cls, name, _capture(getattr(cls, name)))
            # ^no need to juggle with binding the captured method:
            # it will work just as any other method in the class, and
            # `self` will be filled in by the Python runtime itself.
            
            # \/ also, no need to return anything.
        
            
class B(A):
    def foo(self):
        pass