Wrap calls to methods of a Python class

Question

I would like to wrap a number of class methods in Python with the same wrapper.

Conceptually it would look something like this in the simplest scenario:

x = 0 # some arbitrary context

class Base(object):
    def a(self):
       print "a x: %s" % x

    def b(self):
       print "b x: %s" % x

 class MixinWithX(Base):
     """Wrap"""
     def a(self):
         global x
         x = 1
         super(MixinWithX, self).a()
         x = 0

     def b(self):
         global x
         x = 1
         super(MixinWithX, self).a()
         x = 0

Of course, when there are more methods than a and b, this becomes a mess. It seems like there ought to be something simpler. Obviously x could be modified in a decorator but one still ends up having a long list of garbage, which instead of the above looks like:

 from functools import wraps
 def withx(f):
     @wraps(f) # good practice
     def wrapped(*args, **kwargs):
         global x
         x = 1
         f(*args, **kwargs)
         x = 0
     return wrapped

 class MixinWithX(Base):
     """Wrap"""
     @withx
     def a(self):
         super(MixinWithX, self).a()

     @withx
     def b(self):
         super(MixinWithX, self).b()

I thought about using __getattr__ in the mixin, but of course since methods such as a and b are already defined this is never called.

I also thought about using __getattribute__ but it returns the attribute, not wrapping the call. I suppose __getattribute__ could return a closure (example below) but I am not sure how sound a design that is. Here is an example:

 class MixinWithX(Base):
    # a list of the methods of our parent class (Base) that are wrapped
    wrapped = ['a', 'b']

    # application of the wrapper around the methods specified
    def __getattribute__(self, name):
       original = object.__getattribute__(self, name)
       if name in wrapped:
          def wrapped(self, *args, **kwargs):
              global x
              x = 1 # in this example, a context manager would be handy.
              ret = original(*args, **kwargs)
              x = 0
              return ret
          return wrapped
       return original

It has occurred to me that there may be something built into Python that may alleviate the need to manually reproduce every method of the parent class that is to be wrapped. Or maybe a closure in __getattribute__ is the proper way to do this. I would be grateful for thoughts.

Answer 1

Here's my attempt, which allows for a more terse syntax...

x = 0 # some arbitrary context

# Define a simple function to return a wrapped class
def wrap_class(base, towrap):
    class ClassWrapper(base):
        def __getattribute__(self, name):
            original = base.__getattribute__(self, name)
            if name in towrap:
                def func_wrapper(*args, **kwargs):
                    global x
                    x = 1
                    try:
                        return original(*args, **kwargs)
                    finally:
                        x = 0
                return func_wrapper
            return original
    return ClassWrapper


# Our existing base class
class Base(object):
    def a(self):
       print "a x: %s" % x

    def b(self):
       print "b x: %s" % x


# Create a wrapped class in one line, without needing to define a new class
# for each class you want to wrap.
Wrapped = wrap_class(Base, ('a',))

# Now use it
m = Wrapped()
m.a()
m.b()

# ...or do it in one line...
m = wrap_class(Base, ('a',))()

...which outputs...

a x: 1
b x: 0

Answer 2

You can do this using decorators and inspect:

from functools import wraps
import inspect

def withx(f):
    @wraps(f)
    def wrapped(*args, **kwargs):
        print "decorator"
        x = 1
        f(*args, **kwargs)
        x = 0
    return wrapped

class MyDecoratingBaseClass(object):
    def __init__(self, *args, **kwargs):
        for member in inspect.getmembers(self, predicate=inspect.ismethod):
            if member[0] in self.wrapped_methods:
                setattr(self, member[0], withx(member[1]))

class MyDecoratedSubClass(MyDecoratingBaseClass):
    wrapped_methods = ['a', 'b']
    def a(self):
        print 'a'

    def b(self):
        print 'b'

    def c(self):
        print 'c'   

if __name__ == '__main__':
    my_instance = MyDecoratedSubClass()
    my_instance.a()
    my_instance.b()
    my_instance.c()

Output:

decorator
a
decorator
b
c

Answer 3

There are two general directions I can think of which are useful in your case.

One is using a class decorator. Write a function which takes a class, and returns a class with the same set of methods, decorated (either by creating a new class by calling type(...), or by changing the input class in place).

EDIT: (the actual wrapping/inspecting code I had in mind is similar to what @girasquid has in his answer, but connecting is done using decoration instead of mixin/inheritance, which I think is more flexible an robust.)

Which brings me to the second option, which is to use a metaclass, which may be cleaner (yet trickier if you're not used to working with metaclasses). If you don't have access to the definition of the original class, or don't want to change the original definition, you can subclass the original class, setting the metaclass on the derived.