Merge pull request #82 from marcintustin/master

Add full support for Pickle Protocol v2.

Closes #78
Closes #82

Thanks-to: Marcin Tustin <marcin.tustin@gmail.com>
Signed-off-by: David Aguilar <davvid@gmail.com>

jsonpickle/__init__.py

@@ -80,7 +80,8 @@ def encode(value,
            keys=False,
            max_depth=None,
            backend=None,
-           warn=False):
+           warn=False,
+           max_iter=None):
     """
     Return a JSON formatted representation of value, a Python object.
 
@@ -101,6 +102,8 @@ def encode(value,
     :param warn: If set to True then jsonpickle will warn when it
         returns None for an object which it cannot pickle
         (e.g. file descriptors).
+    :param max_iter: If set to a non-negative integer then jsonpickle will
+        consume at most `max_iter` items when pickling iterators.
 
     >>> encode('my string')
     '"my string"'

jsonpickle/pickler.py

@@ -8,7 +8,8 @@
 # you should have received as part of this distribution.
 
 import warnings
-from itertools import chain
+import sys
+from itertools import chain, islice
 
 import jsonpickle.util as util
 import jsonpickle.tags as tags
@@ -26,15 +27,17 @@ def encode(value,
            reset=True,
            backend=None,
            warn=False,
-           context=None):
+           context=None,
+           max_iter=None):
     backend = _make_backend(backend)
     if context is None:
         context = Pickler(unpicklable=unpicklable,
                           make_refs=make_refs,
                           keys=keys,
                           backend=backend,
                           max_depth=max_depth,
-                          warn=warn)
+                          warn=warn,
+                          max_iter=max_iter)
     return backend.encode(context.flatten(value, reset=reset))
 
 
@@ -48,12 +51,13 @@ def _make_backend(backend):
 class Pickler(object):
 
     def __init__(self,
-                unpicklable=True,
-                make_refs=True,
-                max_depth=None,
-                backend=None,
-                keys=False,
-                warn=False):
+                 unpicklable=True,
+                 make_refs=True,
+                 max_depth=None,
+                 backend=None,
+                 keys=False,
+                 warn=False,
+                 max_iter=None):
         self.unpicklable = unpicklable
         self.make_refs = make_refs
         self.backend = _make_backend(backend)
@@ -67,6 +71,8 @@ def __init__(self,
         self._objs = {}
         ## Avoids garbage collection
         self._seen = []
+        # maximum amount of items to take from a pickled iterator
+        self._max_iter = max_iter
 
     def reset(self):
         self._objs = {}
@@ -218,6 +224,8 @@ def _flatten_obj_instance(self, obj):
         has_dict = hasattr(obj, '__dict__')
         has_slots = not has_dict and hasattr(obj, '__slots__')
         has_getnewargs = hasattr(obj, '__getnewargs__')
+        has_getinitargs = hasattr(obj, '__getinitargs__')
+        has_reduce, has_reduce_ex = util.has_reduce(obj)
 
         # Support objects with __getstate__(); this ensures that
         # both __setstate__() and __getstate__() are implemented
@@ -236,14 +244,53 @@ def _flatten_obj_instance(self, obj):
                 data[tags.OBJECT] = class_name
             return handler(self).flatten(obj, data)
 
+        reduce_val = None
         if has_class and not util.is_module(obj):
             if self.unpicklable:
                 class_name = util.importable_name(cls)
                 data[tags.OBJECT] = class_name
 
+            # test for a reduce implementation, and redirect before doing anything else
+            # if that is what reduce requests
+            if has_reduce_ex:
+                try:
+                    # we're implementing protocol 2
+                    reduce_val = obj.__reduce_ex__(2)
+                except TypeError:
+                    # A lot of builtin types have a reduce which just raises a TypeError
+                    # we ignore those
+                    pass
+
+            if has_reduce and not reduce_val:
+                try:
+                    reduce_val = obj.__reduce__()
+                except TypeError:
+                    # A lot of builtin types have a reduce which just raises a TypeError
+                    # we ignore those
+                    pass
+
+            if reduce_val:
+                try:
+                    # At this stage, we only handle the case where __reduce__ returns a string
+                    # other reduce functionality is implemented further down
+                    if isinstance(reduce_val, (str, unicode)):
+                        varpath = iter(reduce_val.split('.'))
+                        # curmod will be transformed by the loop into the value to pickle
+                        curmod = sys.modules[next(varpath)]
+                        for modname in varpath:
+                            curmod = getattr(curmod, modname)
+                            # replace obj with value retrieved
+                        return self._flatten(curmod)
+                except KeyError:
+                    # well, we can't do anything with that, so we ignore it
+                    pass
+
             if has_getnewargs:
                 data[tags.NEWARGS] = self._flatten(obj.__getnewargs__())
 
+            if has_getinitargs:
+                data[tags.INITARGS] = self._flatten(obj.__getinitargs__())
+
         if has_getstate:
             try:
                 state = obj.__getstate__()
@@ -267,6 +314,40 @@ def _flatten_obj_instance(self, obj):
             self._flatten_dict_obj(obj, data)
             return data
 
+        if util.is_sequence_subclass(obj):
+            return self._flatten_sequence_obj(obj, data)
+
+        if util.is_noncomplex(obj):
+            return [self._flatten(v) for v in obj]
+
+        if util.is_iterator(obj):
+            # force list in python 3
+            data[tags.ITERATOR] = list(map(self._flatten, islice(obj, self._max_iter)))
+            return data
+
+        if reduce_val and not isinstance(reduce_val, (str, unicode)):
+            # at this point, reduce_val should be some kind of iterable
+            # pad out to len 5
+            rv_as_list = list(reduce_val)
+            insufficiency = 5 - len(rv_as_list)
+            if insufficiency:
+                rv_as_list+=[None]*insufficiency
+
+            if rv_as_list[0].__name__ == '__newobj__':
+                rv_as_list[0] = tags.NEWOBJ
+
+            data[tags.REDUCE] = list(map(self._flatten, rv_as_list))
+
+            # lift out iterators, so we don't have to iterator and uniterator their content
+            # on unpickle
+            if data[tags.REDUCE][3]:
+                data[tags.REDUCE][3] = data[tags.REDUCE][3][tags.ITERATOR]
+
+            if data[tags.REDUCE][4]:
+                data[tags.REDUCE][4] = data[tags.REDUCE][4][tags.ITERATOR]
+
+            return data
+
         if has_dict:
             # Support objects that subclasses list and set
             if util.is_sequence_subclass(obj):
@@ -276,12 +357,6 @@ def _flatten_obj_instance(self, obj):
             getattr(obj, '_', None)
             return self._flatten_dict_obj(obj.__dict__, data)
 
-        if util.is_sequence_subclass(obj):
-            return self._flatten_sequence_obj(obj, data)
-
-        if util.is_noncomplex(obj):
-            return [self._flatten(v) for v in obj]
-
         if has_slots:
             return self._flatten_newstyle_with_slots(obj, data)
 

jsonpickle/tags.py

@@ -11,23 +11,31 @@
 
 FUNCTION = 'py/function'
 ID = 'py/id'
+INITARGS = 'py/initargs'
+ITERATOR = 'py/iterator'
 JSON_KEY = 'json://'
 NEWARGS = 'py/newargs'
+NEWOBJ = 'py/newobj'
 OBJECT = 'py/object'
-REPR = 'py/repr'
+REDUCE = 'py/reduce'
 REF = 'py/ref'
-STATE = 'py/state'
-SET = 'py/set'
+REPR = 'py/repr'
 SEQ = 'py/seq'
+SET = 'py/set'
+STATE = 'py/state'
 TUPLE = 'py/tuple'
 TYPE = 'py/type'
 
 # All reserved tag names
 RESERVED = set([
     FUNCTION,
     ID,
+    INITARGS,
+    ITERATOR,
     NEWARGS,
+    NEWOBJ,
     OBJECT,
+    REDUCE,
     REF,
     REPR,
     SEQ,

jsonpickle/unpickler.py

@@ -128,10 +128,14 @@ def _restore(self, obj):
             restore = self._restore_id
         elif has_tag(obj, tags.REF): # Backwards compatibility
             restore = self._restore_ref
+        elif has_tag(obj, tags.ITERATOR):
+            restore = self._restore_iterator
         elif has_tag(obj, tags.TYPE):
             restore = self._restore_type
         elif has_tag(obj, tags.REPR): # Backwards compatibility
             restore = self._restore_repr
+        elif has_tag(obj, tags.REDUCE):
+            restore = self._restore_reduce
         elif has_tag(obj, tags.OBJECT):
             restore = self._restore_object
         elif has_tag(obj, tags.FUNCTION):
@@ -148,6 +152,50 @@ def _restore(self, obj):
             restore = lambda x: x
         return restore(obj)
 
+    def _restore_iterator(self, obj):
+        return iter(self._restore_list(obj[tags.ITERATOR]))
+
+    def _restore_reduce(self, obj):
+        """
+        Supports restoring with all elements of __reduce__ as per pep 307.
+        Assumes that iterator items (the last two) are represented as lists
+        as per pickler implementation.
+        """
+        reduce_val = obj[tags.REDUCE]
+        f, args, state, listitems, dictitems = map(self._restore, reduce_val)
+        if f == tags.NEWOBJ or f.__name__ == '__newobj__':
+            # mandated special case
+            cls = args[0]
+            stage1 = cls.__new__(cls, *args[1:])
+        else:
+            stage1 = f(*args)
+
+        if state:
+            try:
+                stage1.__setstate__(state)
+            except AttributeError as err:
+                # it's fine - we'll try the prescribed default methods
+                try:
+                    stage1.__dict__.update(state)
+                except AttributeError as err:
+                    # next prescribed default
+                    for k, v in state.items():
+                        setattr(stage1, k, v)
+
+        if listitems:
+            # should be lists if not None
+            try:
+                stage1.extend(listitems)
+            except AttributeError:
+                for x in listitems:
+                    stage1.append(x)
+
+        if dictitems:
+            for k, v in dictitems:
+                stage1.__setitem__(k, v)
+
+        return stage1
+
     def _restore_id(self, obj):
         return self._objs[obj[tags.ID]]
 
@@ -194,6 +242,7 @@ def _loadfactory(self, obj):
     def _restore_object_instance(self, obj, cls):
         factory = self._loadfactory(obj)
         args = getargs(obj)
+        is_oldstyle = not (isinstance(cls, type) or getattr(cls, '__meta__', None))
 
         # This is a placeholder proxy object which allows child objects to
         # reference the parent object before it has been instantiated.
@@ -203,7 +252,7 @@ def _restore_object_instance(self, obj, cls):
         if args:
             args = self._restore(args)
         try:
-            if hasattr(cls, '__new__'): # new style classes
+            if (not is_oldstyle) and hasattr(cls, '__new__'): # new style classes
                 if factory:
                     instance = cls.__new__(cls, factory, *args)
                     instance.default_factory = factory
@@ -212,10 +261,16 @@ def _restore_object_instance(self, obj, cls):
             else:
                 instance = object.__new__(cls)
         except TypeError: # old-style classes
+            is_oldstyle = True
+
+        if is_oldstyle:
             try:
-                instance = cls()
-            except TypeError: # fail gracefully
-                return self._mkref(obj)
+                instance = cls(*args)
+            except TypeError:  # fail gracefully
+                try:
+                    instance = make_blank_classic(cls)
+                except:  # fail gracefully
+                    return self._mkref(obj)
 
         proxy.instance = instance
         self._swapref(proxy, instance)
@@ -420,6 +475,10 @@ def getargs(obj):
     # Let saved newargs take precedence over everything
     if has_tag(obj, tags.NEWARGS):
         return obj[tags.NEWARGS]
+
+    if has_tag(obj, tags.INITARGS):
+        return obj[tags.INITARGS]
+
     try:
         seq_list = obj[tags.SEQ]
         obj_dict = obj[tags.OBJECT]
@@ -434,6 +493,21 @@ def getargs(obj):
     return []
 
 
+class _trivialclassic:
+    """
+    A trivial class that can be instantiated with no args
+    """
+
+def make_blank_classic(cls):
+    """
+    Implement the mandated strategy for dealing with classic classes
+    which cannot be instantiated without __getinitargs__ because they
+    take parameters
+    """
+    instance = _trivialclassic()
+    instance.__class__ = cls
+    return instance
+
 def loadrepr(reprstr):
     """Returns an instance of the object from the object's repr() string.
     It involves the dynamic specification of code.