16

I want to create a class that wraps another class so that when a function is run through the wrapper class a pre and post function is run as well. I want the wrapper class to work with any class without modification.

For example if i have this class.

class Simple(object):
    def one(self):
        print "one"

    def two(self,two):
        print "two" + two

    def three(self):
        print "three"

I could use it like this...

number = Simple()
number.one()
number.two("2")

I have so far written this wrapper class...

class Wrapper(object):
    def __init__(self,wrapped_class):
        self.wrapped_class = wrapped_class()

    def __getattr__(self,attr):
        return self.wrapped_class.__getattribute__(attr)

    def pre():
        print "pre"

    def post():
        print "post"

Which I can call like this...

number = Wrapper(Simple)
number.one()
number.two("2")

Which can be used the same as above apart from changing the first line.

What I want to happen is when calling a function through the wrapper class the pre function in the wrapper class gets called then the desired function in the wrapped class then the post function. I want to be able to do this without changing the wrapped class and also without changing the way the functions are called, only changing the syntax of how the instance of the class is created. eg number = Simple() vs number = Wrapper(Simple)

SilentGhost
  • 307,395
  • 66
  • 306
  • 293

2 Answers2

31

You're almost there, you just need to do some introspection inside __getattr__, returning a new wrapped function when the original attribute is callable:

class Wrapper(object):
    def __init__(self,wrapped_class):
        self.wrapped_class = wrapped_class()

    def __getattr__(self,attr):
        orig_attr = self.wrapped_class.__getattribute__(attr)
        if callable(orig_attr):
            def hooked(*args, **kwargs):
                self.pre()
                result = orig_attr(*args, **kwargs)
                # prevent wrapped_class from becoming unwrapped
                if result == self.wrapped_class:
                    return self
                self.post()
                return result
            return hooked
        else:
            return orig_attr

    def pre(self):
        print ">> pre"

    def post(self):
        print "<< post"

Now with this code:

number = Wrapper(Simple)

print "\nCalling wrapped 'one':"
number.one()

print "\nCalling wrapped 'two':"
number.two("2")

The result is:

Calling wrapped 'one':
>> pre
one
<< post

Calling wrapped 'two':
>> pre
two2
<< post
Matti Lyra
  • 12,828
  • 8
  • 49
  • 67
Steven Kryskalla
  • 14,179
  • 2
  • 40
  • 42
  • 4
    This is a nice answer, but it doesn't work if the wrapped class has methods that return another instance of the wrapped class. This is not hypothetical - for instance, it's why the pandas package is so hard to subclass, as I have learned in the last couple days :) – Dave Kielpinski Aug 10 '17 at 23:24
  • Doesn't work for inner class calls either - this only wraps calls from outside – D.Ginzbourg Sep 26 '18 at 01:45
3

I have just noticed in my original design there is no way of passing args and kwargs to the wrapped class, here is the answer updated to pass the inputs to the wrapped function...

class Wrapper(object):
def __init__(self,wrapped_class,*args,**kargs):
    self.wrapped_class = wrapped_class(*args,**kargs)

def __getattr__(self,attr):
    orig_attr = self.wrapped_class.__getattribute__(attr)
    if callable(orig_attr):
        def hooked(*args, **kwargs):
            self.pre()
            result = orig_attr(*args, **kwargs)
            self.post()
            return result
        return hooked
    else:
        return orig_attr

def pre(self):
    print ">> pre"

def post(self):
    print "<< post"