Python Wrap Class Method

Joe J picture Joe J · Jul 21, 2011 · Viewed 14.3k times · Source

I'm trying to create an object with a run method that will be wrapped by a _wrap_run method. I'd like to be able to call the method and it's wrapper by simply typing instance.run() and I'd like to be able to subclass the object so I can override the run() method and have it still execute the wrapper.

More simply put, I want people to be able to subclass A and override run() but still have calls to the run() method execute the wrapper function.

I'm having some difficulty with the mechanics of this. Does anyone have any suggestions regarding this approach?

class A:

    def run(self):
        print "Run A"
        return True

    def _wrap_run(self):
        print "PRE"
        return_value = self.run()
        print "POST"
        return return_value

    run = property(_wrap_run)


a = A()
a.run()
"""
Should Print: 
PRE
Run A
POST
"""


class B(A):

    def run(self):
        print "Run B"
        return True

b = B()
b.run()
"""
Should Print: 
PRE
Run B
POST
"""

Answer

agf picture agf · Jul 21, 2011

Use a Metaclass.

class MetaClass(type):
    @staticmethod
    def wrap(run):
        """Return a wrapped instance method"""
        def outer(self):
            print "PRE",
            return_value = run(self)
            print "POST"
            return return_value
        return outer
    def __new__(cls, name, bases, attrs):
        """If the class has a 'run' method, wrap it"""
        if 'run' in attrs:
            attrs['run'] = cls.wrap(attrs['run'])
        return super(MetaClass, cls).__new__(cls, name, bases, attrs)

class MyClass(object):
    """Use MetaClass to make this class"""
    __metaclass__ = MetaClass
    def run(self): print 'RUN',

myinstance = MyClass()

# Prints PRE RUN POST
myinstance.run()

Now if other people subclass MyClass, they will still get their run() methods wrapped.