Wrapping a class whose methods return instances of that class

Viewed 642

I need to write a class to wrap classes from third-party packages. Usually, the third-party class has methods that return third-party class instances. The wrapped versions of these methods have to convert those instances into instances of the wrapped class, but I can't make it work. I'm using Python 2.7 with new-style classes.

Based on Create a wrapper class to call a pre and post function around existing functions?, I've got the following.

import copy

class Wrapper(object):
    __wraps__  = None

    def __init__(self, obj):
        if self.__wraps__ is None:
            raise TypeError("base class Wrapper may not be instantiated")
        elif isinstance(obj, self.__wraps__):
            self._obj = obj
        else:
            raise ValueError("wrapped object must be of %s" % self.__wraps__)

    def __getattr__(self, name):
        orig_attr = self._obj.__getattribute__(name)
        if callable(orig_attr):
            def hooked(*args, **kwargs):
                result = orig_attr(*args, **kwargs)
                if result == self._obj:
                    return result
                return self.__class__(result)
            return hooked
        else:
            return orig_attr

class ClassToWrap(object):
    def __init__(self, data):
        self.data = data

    def theirfun(self):
        new_obj = copy.deepcopy(self)
        new_obj.data += 1
        return new_obj

class Wrapped(Wrapper):
    __wraps__ = ClassToWrap

    def myfun(self):
        new_obj = copy.deepcopy(self)
        new_obj.data += 1
        return new_obj

obj = ClassToWrap(0)
wr0 = Wrapped(obj)
print wr0.data
>> 0
wr1 = wr0.theirfun()
print wr1.data
>> 1
wr2 = wr1.myfun()
print wr2.data
>> 2
wr3 = wr2.theirfun()
print wr3.data
>> 2

Now why on earth does theirfun() work the first time, but not the second time? Both wr0 and wr2 are of type Wrapped, and invoking wr2.theirfun() doesn't raise an error, but it doesn't add 1 to wr2.data as expected.

Sorry, but I am not looking for the following alternative approaches:

  1. Monkey patching. My codebase is nontrivial and I don't know how to ensure the patch will propagate through the web of import statements.
  2. Writing individual wrapper methods for all these tricky methods for each third-party package. There are too many of them.

ETA: There are a couple helpful answers that reference the underlying _obj attribute outside of the Wrapper class. However, the point of this approach is extensibility, so this functionality needs to be inside the Wrapper class. myfun needs to behave as expected without referencing _obj in its definition.

2 Answers
Related