Search code examples
pythonwrappermetaclass

How to wrap every method of a class?


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?


Solution

  • 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