How to wrap every method of a class? [duplicate]

Question

I'd like to wrap every method of a particular class in python, and I'd like to do so by editing the code of the class minimally. How should I go about this?

Christoperchristoph

Answer 1

An elegant way to do it is described in Michael Foord's Voidspace blog in an entry about what metaclasses are and how to use them in the section titled A Method Decorating Metaclass. Simplifying it slightly and applying it to your situation resulted in this:

from functools import wraps
from types import FunctionType

def wrapper(method):
    @wraps(method)
    def wrapped(*args, **kwargs):
    #   ... <do something to/with "method" or the result of calling it>
    return wrapped

class MetaClass(type):
    def __new__(meta, classname, bases, classDict):
        newClassDict = {}
        for attributeName, attribute in classDict.items():
            if isinstance(attribute, FunctionType):
                # replace it with a wrapped version
                attribute = wrapper(attribute)
            newClassDict[attributeName] = attribute
        return type.__new__(meta, classname, bases, newClassDict)

class MyClass(object):
    __metaclass__ = MetaClass  # wrap all the methods
    def method1(self, ...):
        # ...etc ...

In Python, function/method decorators are just function wrappers plus some syntactic sugar to make using them easy (and prettier).

Python 3 Compatibility Update

The previous code uses Python 2.x metaclass syntax which would need to be translated in order to be used in Python 3.x, however it would then no longer work in the previous version. This means it would need to use:

class MyClass(metaclass=MetaClass)  # apply method-wrapping metaclass
    ...

instead of:

class MyClass(object):
    __metaclass__ = MetaClass  # wrap all the methods
    ...

If desired, it's possible to write code which is compatible with both Python 2.x and 3.x, but doing so requires using a slightly more complicated technique which dynamically creates a new base class that inherits the desired metaclass, thereby avoiding errors due to the syntax differences between the two versions of Python. This is basically what Benjamin Peterson's six module's with_metaclass() function does.

from types import FunctionType
from functools import wraps

def wrapper(method):
    @wraps(method)
    def wrapped(*args, **kwargs):
        print('{!r} executing'.format(method.__name__))
        return method(*args, **kwargs)
    return wrapped


class MetaClass(type):
    def __new__(meta, classname, bases, classDict):
        newClassDict = {}
        for attributeName, attribute in classDict.items():
            if isinstance(attribute, FunctionType):
                # replace it with a wrapped version
                attribute = wrapper(attribute)
            newClassDict[attributeName] = attribute
        return type.__new__(meta, classname, bases, newClassDict)


def with_metaclass(meta):
    """ Create an empty class with the supplied bases and metaclass. """
    return type.__new__(meta, "TempBaseClass", (object,), {})


if __name__ == '__main__':

    # Inherit metaclass from a dynamically-created base class.
    class MyClass(with_metaclass(MetaClass)):
        @staticmethod
        def a_static_method():
            pass

        @classmethod
        def a_class_method(cls):
            pass

        def a_method(self):
            pass

    instance = MyClass()
    instance.a_static_method()  # Not decorated.
    instance.a_class_method()   # Not decorated.
    instance.a_method()         # -> 'a_method' executing

Beitch

BTW, recently I saw a very comprehensive answer to the question How to make built-in containers (sets, dicts, lists) thread safe? which describes many different ways to wrap methods. I think you might find it very interesting. – Beitch

Thank you for this @martineau. Would it be possible to demonstrate the Python3 version of your answer with a concrete example. Reason: I'm having trouble deciphering which of the above are key-words, and which are meant for substitution when actually applying it. I'd love to try this with, say, the 'pandas.DataFrame' class. :) – Selfreliance

@prismalytics: Sure. See runnable demo I created that works unchanged in both Python 2 & 3: wrapping_methods_example.py – Beitch

Thank you @martineau. I will study the example you pointed me to and apply it to my use case. Very much appreciated (with upvotes all over the show it). =:) – Selfreliance

@prismalytics: You're welcome, but up-voting all over isn't necessary — besides the system may notice the pattern and undo them. Anyway, there's a lot going on this answer's code, so if you have a specific question, just ask (or even post a new question). – Beitch

Hi @martineau. I spent yesterday studying your above paste example (thanks). I'm having a little trouble adapting it. Here is a paste of my adaptation. Notice the second call to .info() at the bottom. I expected it to print out extra information (defined in wrapper()), but it doesn't. Any idea why? – Selfreliance

P.S. @martineau, I noticed and fixed a bug in the paste, but it still doesn't behave as I thought it would. – Selfreliance

@Beitch Comparing types using type() and == is not recommended - use isinstance(attribute, types.FunctionType) instead! docs.quantifiedcode.com/python-code-patterns/readability/… – Wham

@Monkpit: Point taken. It's there because I just copied some of Michael Foord's code — fixed — although it's hard to imagine anyone ever subclassing types.FunctionType.... – Beitch

For anyone following along, @prismalytics.io asked a new but related question titled Wrapping all possible method calls of a class in a try/except block. – Beitch

Love this! So much cleaner than overriding __getattribute__, and easier to document for the user. – Thedrick

@LukeDavis: How about showing your love with an upvote? ;¬) – Beitch

Answer 2

You mean programatically set a wrapper to methods of a class?? Well, this is probably a really bad practice, but here's how you may do it:

def wrap_methods( cls, wrapper ):
    for key, value in cls.__dict__.items( ):
        if hasattr( value, '__call__' ):
            setattr( cls, key, wrapper( value ) )

If you have class, for example

class Test( ):
    def fire( self ):
        return True
    def fire2( self ):
        return True

and a wrapper

def wrapper( fn ):
    def result( *args, **kwargs ):
        print 'TEST'
        return fn( *args, **kwargs )
    return result

then calling

wrap_methods( Test, wrapper )

will apply wrapper to all methods defined in class Test. Use with caution! Actually, don't use it at all!

Banting

I don't intend to build with it -- it's just a debugging tool that I want. Thanks! – Christoperchristoph

Decorating result function with @wraps(fn) yields more convenient state (setting method name, etc.) - see docs.python.org/2/library/functools.html#functools.wraps – Endermic

Answer 3

If extensively modifying default class behavior is the requirement, MetaClasses are the way to go. Here's an alternative approach.

If your use case is limited to just wrapping instance methods of a class, you could try overriding the __getattribute__ magic method.

from functools import wraps
def wrapper(func):
    @wraps(func)
    def wrapped(*args, **kwargs):
        print "Inside Wrapper. calling method %s now..."%(func.__name__)
        return func(*args, **kwargs)
    return wrapped

Make sure to use functools.wraps while creating wrappers, even more so if the wrapper is meant for debugging since it provides sensible TraceBacks.

import types
class MyClass(object): # works only for new-style classes
    def method1(self):
        return "Inside method1"
    def __getattribute__(self, name):
        attr = super(MyClass, self).__getattribute__(name)
        if type(attr) == types.MethodType:
            attr = wrapper(attr)
        return attr

Malamud

I think it's worth pointing out that this approach (re)wraps all methods every time they're called, which entails significantly more overhead than if the wrapping was just done once and made a part of the class, as can be done with a metaclass or a class decorator. Of course this additional overhead might be perfectly acceptable if it's only being done for debugging purposes. – Beitch

@martineau: Very valid point. I also should have mentioned that I always shied away from MetaClasses (it seems like very fragile space to me), up until now. – Malamud