changeset 717:fcfbe103cfd1

internal: Fix some bugs with the `fastpickle` module.
author Ludovic Chabant <ludovic@chabant.com>
date Sat, 28 May 2016 15:41:05 -0700
parents 42efc3634763
children d4408fbbbc7d
files piecrust/fastpickle.py tests/test_fastpickle.py
diffstat 2 files changed, 29 insertions(+), 7 deletions(-) [+]
line wrap: on
line diff
--- a/piecrust/fastpickle.py	Thu May 26 20:45:34 2016 -0700
+++ b/piecrust/fastpickle.py	Sat May 28 15:41:05 2016 -0700
@@ -11,6 +11,12 @@
     return data.encode('utf8')
 
 
+def pickle_obj(obj):
+    if obj is not None:
+        return _pickle_object(obj)
+    return None
+
+
 def pickle_intob(obj, buf):
     data = _pickle_object(obj)
     buf = _WriteWrapper(buf)
@@ -22,6 +28,12 @@
     return _unpickle_object(data)
 
 
+def unpickle_obj(data):
+    if data is not None:
+        return _unpickle_object(data)
+    return None
+
+
 def unpickle_fromb(buf, bufsize):
     with buf.getbuffer() as innerbuf:
         data = codecs.decode(innerbuf[:bufsize], 'utf8')
@@ -70,7 +82,8 @@
     elif op == _UNPICKLING:
         res = collections.OrderedDict()
         for k, v in obj.items():
-            res[k] = func(v)
+            if k != '__type__':
+                res[k] = func(v)
         return res
 
 
@@ -230,13 +243,11 @@
     class_def = getattr(mod, class_name)
     obj = class_def.__new__(class_def)
 
-    del state['__class__']
-    del state['__module__']
     attr_names = list(state.keys())
     for name in attr_names:
-        state[name] = _unpickle_object(state[name])
-
-    obj.__dict__.update(state)
+        if name == '__class__' or name == '__module__':
+            continue
+        obj.__dict__[name] = _unpickle_object(state[name])
 
     return obj
 
--- a/tests/test_fastpickle.py	Thu May 26 20:45:34 2016 -0700
+++ b/tests/test_fastpickle.py	Sat May 28 15:41:05 2016 -0700
@@ -1,6 +1,6 @@
 import datetime
 import pytest
-from piecrust.fastpickle import pickle, unpickle
+from piecrust.fastpickle import pickle, unpickle, pickle_obj, unpickle_obj
 
 
 class Foo(object):
@@ -51,3 +51,14 @@
     for i in range(2):
         assert f.bars[i].value == o.bars[i].value
 
+
+def test_reentrance():
+    a = {'test_ints': 42, 'test_set': set([1, 2])}
+    data = pickle_obj(a)
+    b = unpickle_obj(data)
+    assert a == b
+    other_b = unpickle_obj(data)
+    assert a == other_b
+    c = unpickle_obj(data)
+    assert a == c
+