changeset 450:298f8f46432a

internal: Add a `fastpickle` module to help with multiprocess serialization. Looks like we're wasting a lot of time serializing custom class information for the workers, so this new module helps to convert classes to standard structures (dicts, lists, etc).
author Ludovic Chabant <ludovic@chabant.com>
date Mon, 06 Jul 2015 21:29:17 -0700
parents 30f2c2a595f5
children 838f3964f400
files piecrust/fastpickle.py tests/test_fastpickle.py
diffstat 2 files changed, 167 insertions(+), 0 deletions(-) [+]
line wrap: on
line diff
--- /dev/null	Thu Jan 01 00:00:00 1970 +0000
+++ b/piecrust/fastpickle.py	Mon Jul 06 21:29:17 2015 -0700
@@ -0,0 +1,114 @@
+import sys
+import datetime
+import collections
+
+
+def pickle(obj):
+    return _pickle_object(obj)
+
+
+def unpickle(data):
+    return _unpickle_object(data)
+
+
+def _tuple_dispatch(obj, func):
+    res = [None] * len(obj)
+    for i, c in enumerate(obj):
+        res[i] = func(c)
+    return tuple(res)
+
+
+def _list_dispatch(obj, func):
+    res = [None] * len(obj)
+    for i, c in enumerate(obj):
+        res[i] = func(c)
+    return res
+
+
+def _dict_dispatch(obj, func):
+    res = {}
+    for k, v in obj.items():
+        res[k] = func(v)
+    return res
+
+
+def _set_dispatch(obj, func):
+    res = set()
+    for v in obj:
+        res.add(func(v))
+    return res
+
+
+_identity_dispatch = object()
+
+_type_dispatch = {
+        type(None): _identity_dispatch,
+        bool: _identity_dispatch,
+        int: _identity_dispatch,
+        float: _identity_dispatch,
+        str: _identity_dispatch,
+        datetime.date: _identity_dispatch,
+        datetime.datetime: _identity_dispatch,
+        datetime.time: _identity_dispatch,
+        tuple: _tuple_dispatch,
+        list: _list_dispatch,
+        dict: _dict_dispatch,
+        collections.OrderedDict: _dict_dispatch,
+        set: _set_dispatch
+        }
+
+
+def _pickle_object(obj):
+    t = type(obj)
+    disp = _type_dispatch.get(t)
+    if disp is _identity_dispatch:
+        return obj
+
+    if disp is not None:
+        return disp(obj, _pickle_object)
+
+    if isinstance(obj, Exception):
+        return obj
+
+    getter = getattr(obj, '__getstate__', None)
+    if getter is not None:
+        state = getter()
+    else:
+        state = obj.__dict__
+
+    state = _dict_dispatch(state, _pickle_object)
+    state['__class__'] = obj.__class__.__name__
+    state['__module__'] = obj.__class__.__module__
+
+    return state
+
+
+def _unpickle_object(state):
+    t = type(state)
+    disp = _type_dispatch.get(t)
+    if disp is _identity_dispatch:
+        return state
+
+    if (disp is not None and
+            (t != dict or '__module__' not in state)):
+        return disp(state, _unpickle_object)
+
+    if isinstance(state, Exception):
+        return state
+
+    mod_name = state['__module__']
+    mod = sys.modules[mod_name]
+    class_name = state['__class__']
+    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)
+
+    return obj
+
--- /dev/null	Thu Jan 01 00:00:00 1970 +0000
+++ b/tests/test_fastpickle.py	Mon Jul 06 21:29:17 2015 -0700
@@ -0,0 +1,53 @@
+import datetime
+import pytest
+from piecrust.fastpickle import pickle, unpickle
+
+
+class Foo(object):
+    def __init__(self, name):
+        self.name = name
+        self.bars = []
+
+
+class Bar(object):
+    def __init__(self, value):
+        self.value = value
+
+
+@pytest.mark.parametrize(
+        'obj, expected',
+        [
+            (True, True),
+            (42, 42),
+            (3.14, 3.14),
+            (datetime.date(2015, 5, 21), datetime.date(2015, 5, 21)),
+            (datetime.datetime(2015, 5, 21, 12, 55, 32),
+                datetime.datetime(2015, 5, 21, 12, 55, 32)),
+            (datetime.time(9, 25, 57), datetime.time(9, 25, 57)),
+            ((1, 2, 3), (1, 2, 3)),
+            ([1, 2, 3], [1, 2, 3]),
+            ({'foo': 1, 'bar': 2}, {'foo': 1, 'bar': 2}),
+            (set([1, 2, 3]), set([1, 2, 3])),
+            ({'foo': [1, 2, 3], 'bar': {'one': 1, 'two': 2}},
+                {'foo': [1, 2, 3], 'bar': {'one': 1, 'two': 2}})
+            ])
+def test_pickle_unpickle(obj, expected):
+    data = pickle(obj)
+    actual = unpickle(data)
+    assert actual == expected
+
+
+def test_objects():
+    f = Foo('foo')
+    f.bars.append(Bar(1))
+    f.bars.append(Bar(2))
+
+    data = pickle(f)
+    o = unpickle(data)
+
+    assert type(o) == Foo
+    assert o.name == 'foo'
+    assert len(o.bars) == 2
+    for i in range(2):
+        assert f.bars[i].value == o.bars[i].value
+