Mercurial > piecrust2
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 +