jsonpickle_convert_to_setstate.py

#!/usr/bin/env python

# jsonpickle doesn't handle the case of an object being serialized without
# __getstate__ to then be restored with __setstate__. __setstate__ is never
# called because the json data does not have the "py/state" key. Only by adding
# a custom handler can we support both formats to go through __setstate__

import jsonpickle

class Original(object):
    def __init__(self, a, b):
        self.a = a
        self.b = b

items = [Original(1,5), Original(2,"stuff")]

class AddedAttribute(object):
    def __init__(self, a, b):
        print("AddedAttribute.__init__!!!!")
        self.a = a
        self.b = b

    def __repr__(self):
        return ",".join(sorted(["%s:%s" % a for a in self.__dict__.items()]))

    def __getstate__(self):
        return self.__dict__

    def __setstate__(self, state):
        print("setstate!", state)
        self.extra = state.pop("EXTRA!", None)
        self.__dict__.update(state)

items = [Original(1,5), Original(2,"stuff")]

encoded_original = jsonpickle.encode(items)
print(encoded_original)
# produces: [{"py/object": "__main__.AddedAttribute", "py/state": {"a": 1, "b": 5}}, {"py/object": "__main__.AddedAttribute", "py/state": {"a": 2, "b": "stuff"}}]

items = [AddedAttribute(1,5), AddedAttribute(2,"stuff")]

encoded_with_added = jsonpickle.encode(items)
print(encoded_with_added)
# produces: [{"py/object": "__main__.AddedAttribute", "py/state": {"a": 1, "b": 5}}, {"py/object": "__main__.AddedAttribute", "py/state": {"a": 2, "b": "stuff"}}]
#
# Note the py/state key! The presence of this triggers the use of __setstate__.
# Without the json in this format, __setstate__ is not called.

decoded_with_added = jsonpickle.decode(encoded_with_added)
print(decoded_with_added)
# produces: [a:1,b:5,extra:None, a:2,b:stuff,extra:None]

# Transform Original into AddedAttribute with simple string replacement to test
# if setstate is called
modified_encoded_original = encoded_original.replace("Original", "AddedAttribute")
modified_decoded_original = jsonpickle.decode(modified_encoded_original)
print(modified_decoded_original)
# produces: [a:1,b:5, a:2,b:stuff]

# Only by adding a custom handler can we support both formats to go through
# __setstate__
class ExtraAttributeHandler(jsonpickle.handlers.BaseHandler):
    def flatten(self, obj, data):
        print("FLATTEN!")
        print(obj)
        print(data)
        data["py/state"] = obj.__getstate__()
        return data

    def restore(self, data):
        print("RESTORE!")
        print(data)
        print(self.context)
        cls = jsonpickle.unpickler.loadclass(data["py/object"])
        print(cls)
        restored = cls.__new__(cls)
        if "py/state" in data:
            state = data["py/state"]
        else:
            state = {k:v for k,v in data.iteritems() if not k.startswith("py/")}
        restored.__setstate__(state)
        return restored

jsonpickle.handlers.register(AddedAttribute, ExtraAttributeHandler)
extra_attribute_original = jsonpickle.decode(modified_encoded_original)
print(extra_attribute_original)
# produces: [a:1,b:5,extra:None, a:2,b:stuff,extra:None]

e = jsonpickle.encode(extra_attribute_original)
print(e)
u = jsonpickle.decode(e)
print(u)