# All instances are stored in this dictionary.
instance_dict = {}

# A persistent object class, whereby the pid value given is the
# reference we use for the object.
class PObject:
    def __init__(self, pid):
        self.pid = pid
        instance_dict[pid] = self
    def __repr__(self):
        return '%s(%s)' % (self.__class__.__name__, repr(self.pid))

# Makes a persistent reference for PObject instances.
def persistent_id(obj):
    if isinstance(obj, PObject):
        return obj.pid
    return None

# Handles references for PObject instances.
def persistent_load(pid):
    if instance_dict.has_key(pid):
        return instance_dict[pid]
    return PObject(pid)

def test(instance, pickle, bin):
    import StringIO
    s = StringIO.StringIO()

    # Create a pickler from the provided pickling module,
    # and override the implementation of persistent_id
    # given.
    pickler = pickle.Pickler(s, bin)
    pickler.persistent_id = persistent_id
    pickler.dump(instance)

    # Grab the pickled string (referenced in pickled).
    s.seek(0)
    pickled = s.read()
    s.seek(0)

    # Now attempt to unpickle the data we have just pickled,
    # using the version of persistent_load defined at the
    # module level.
    unpickler = pickle.Unpickler(s)
    unpickler.persistent_load = persistent_load
    try:
        loaded = unpickler.load()
    except (Exception, pickle.UnpicklingError), exc:
        success = 0
    else:
        success = loaded is instance
        exc = None

    return pickled, success, exc, s.read()


if __name__ == '__main__':

    import cPickle
    import pickle
    import marshal

    instance_okay = PObject('spam')
    instance_bad = PObject('two\nlines')
    instance_bad_2 = PObject(marshal.dumps('ten chars!'))

    for instance in [instance_okay, instance_bad, instance_bad_2]:
        for bin in [1, 0]:
            for module in [pickle, cPickle]:
                print "Module: %s    Instance: %s  Bin=%s" % (module.__name__, instance, bin)
                pickled, success, exc, unread = test(instance, module, bin)
                if success:
                    print "Successful round trip."
                elif exc is None:
                    print "Failed round trip (didn't get the right instance back)"
                else:
                    print "Failed round trip: %s" % (exc,)
                print "Pickled as: %s" % (repr(pickled),)
                if unread:
                    print "Unread data: %s" % (repr(unread),)
                print