Mercurial > piecrust2
diff piecrust/fastpickle.py @ 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 | |
children | 2ef04e16f0b9 |
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 +