wujek
wujek

Reputation: 11040

Patching all instances of a class in Python

I'm on Python 3.5.1. I wrote a library, and under certain circumstances (when the file is run directly, i.e. __name__ == '__main__') I want to decorate certain methods in one of the classes. It should decorate all the instances that can ever be created. I would like to do it in a way that is non invasive, i.e. ideally the class in my library would not need any special code.

After a while, I managed to implement something like this, which meets my requirements:

def patch(clazz, name, replacement):

    def wrap_original(orig):
        # when called with the original function, a new function will be returned
        # this new function, the wrapper, replaces the original function in the class
        # and when called it will call the provided replacement function with the
        # original function as first argument and the remaining arguments filled in by Python

        def wrapper(*args, **kwargs):
            return replacement(orig, *args, **kwargs)

        return wrapper

    orig = getattr(clazz, name)
    setattr(clazz, name, wrap_original(orig))

def replacement_function(orig, self, ... other argumnents ...):
    # orig here is the original function, so can be called like this:
    # orig(self, ... args ...)
    pass

patch(mylib.MyClass, 'method_name', replacemment_function)

Amazingly, this code works, although I haven't tested it with class methods yet, but I don't need it now. It also patches instances created before the patching, although I'm not yet certain whether it's good or not ;d

The code above is arguably difficult, I needed a while to wrap my head around the way it works after I wrote it, in order to write the explanation comment. I would love something easier.

The question: is there anything in the Python library that would make code like this unnecessary, which already implements what I'm doing, but better?

Upvotes: 3

Views: 979

Answers (4)

wujek
wujek

Reputation: 11040

One of the posters here, who sadly deleted her/his post, directed me towards the functools module. In the end, I settled on the following:

def replacement(self, orig, ... other arguments ...):
    # orig here is the original function, so can be called like this:
    # orig(self, ... args ...)
    pass

mylib.MyClass.my_method = functools.partialmethod(replacement, mylib.MyClass.my_method)

The orig and self arguments needed to switch places, as the result of partialmethod binds the first argument to the instance it is in, and the second in this case will be the original function (the second argument to partialmethod). Looks much cleaner.

Upvotes: 0

aghast
aghast

Reputation: 15310

Another alternative would be to create a "null" decorator function, and then switch between that function and the "real" decorator using your conditional logic:

from decorator_lib import real_decorator

def fake_decorator(fun):
    return fun

if __name__ == '__main__':
    my_decorator = real_decorator
else:
    my_decorator = fake_decorator


# ... elsewhere in the module ...

@my_decorator
def method(self, a, b, c):
    pass

# ... finally:
if __name__ == '__main__':
    run_self_tests_or_whatever()

Upvotes: 0

Martijn Pieters
Martijn Pieters

Reputation: 1121834

Methods are created dynamically when looked up on instances; instances do not have copies of all methods, instead the descriptor protocol takes functions from the class and binds those to the instance as needed. This is why monkeypatching the class works here; instance.method_name will find mylib.MyClass.method_name when the attribute lookup is performed.

There is nothing in the default library that performs what you are doing here, no, because different code may need different patterns of handling delegation back to the old method.

Your approach looks very close to how the Mercurial project supports function wrapping, in that the original is passed into the wrapper.

Upvotes: 3

Rushy Panchal
Rushy Panchal

Reputation: 17532

Your approach seems to be the most Pythonic way to go about it.

Gevent, a popular library that makes use of monkey patching, performs monkey patching in almost the same way you describe.

Upvotes: 1

Related Questions