Wrapping a class whose methods return instances of that class

Question

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.

Answer 1

The problem is with your implementation of myfun in the Wrapped class. You only update the data member of the class' instance, but the wrapped class' (ClassToWrap instance i.e. _obj) data member is obsolete, using a value from the prev call of theirfun.

You need to synchronise the data values across both instances:

class Wrapper(object):
    ...
    def __setattr__(self, attr, val):
        object.__setattr__(self, attr, val)
        if getattr(self._obj, attr, self._obj) is not self._obj: # update _obj's member if it exists
            setattr(self._obj, attr, getattr(self, attr))


class Wrapped(Wrapper):
    ...
    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
# 3
wr4 = wr3.myfun()
print wr4.data
# 4
wr5 = wr4.theirfun()
print wr5.data
# 5

Answer 2

The issue is with the assignment new_obj.data += 1 in myfun. The problem with it is that new_obj is an instance of Wrapped, not an instance of ClassToWrap. Your Wrapper base class only supports looking up attributes on the proxied object. It doesn't support assignments to attributes. Augmented assignment does both, so it doesn't work entirely correctly.

You could make your myfun work by changing it around a little bit:

def myfun(self):
    new_obj = copy.deepcopy(self._obj)
    new_obj.data += 1
    return self.__class__(new_obj) # alternative spelling: return type(self)(new_obj)

Another approach to solving the issue would be to add a __setattr__ method to Wrapper, but getting it to work right (without interfering in the proxy class's own attributes) would be a bit awkward.

Unrelated to your current issue, you're also potentially leaking the wrapped object in the hooked wrapper function. If a method you've called returns the object it was called upon (e.g. the method did return self), you're currently returning that object unwrapped. You probably want to change return result to return self so that you return the current wrapper. You may also want a check on the return value to see if it is a type you're able to wrap or not. Currently your code will fail if a method returns a string or number or anything other than an instance of the wrapped type.

        def hooked(*args, **kwargs):
            result = orig_attr(*args, **kwargs)
            if result == self._obj:
                return self                            # fix for leaking objects
            elif isisntance(result, self.__wraps__): # new type check
                return self.__class__(result)
            else:
                return result              # fallback for non-wrappable return values