Mercurial > piecrust2

comparison piecrust/fastpickle.py @ 450:298f8f46432a

Find changesets by keywords (author, files, the commit message), revision number or hash, or revset expression.
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
comparison
equal deleted inserted replaced
449:30f2c2a595f5 450:298f8f46432a
1 import sys
2 import datetime
3 import collections
4
5
6 def pickle(obj):
7 return _pickle_object(obj)
8
9
10 def unpickle(data):
11 return _unpickle_object(data)
12
13
14 def _tuple_dispatch(obj, func):
15 res = [None] * len(obj)
16 for i, c in enumerate(obj):
17 res[i] = func(c)
18 return tuple(res)
19
20
21 def _list_dispatch(obj, func):
22 res = [None] * len(obj)
23 for i, c in enumerate(obj):
24 res[i] = func(c)
25 return res
26
27
28 def _dict_dispatch(obj, func):
29 res = {}
30 for k, v in obj.items():
31 res[k] = func(v)
32 return res
33
34
35 def _set_dispatch(obj, func):
36 res = set()
37 for v in obj:
38 res.add(func(v))
39 return res
40
41
42 _identity_dispatch = object()
43
44 _type_dispatch = {
45 type(None): _identity_dispatch,
46 bool: _identity_dispatch,
47 int: _identity_dispatch,
48 float: _identity_dispatch,
49 str: _identity_dispatch,
50 datetime.date: _identity_dispatch,
51 datetime.datetime: _identity_dispatch,
52 datetime.time: _identity_dispatch,
53 tuple: _tuple_dispatch,
54 list: _list_dispatch,
55 dict: _dict_dispatch,
56 collections.OrderedDict: _dict_dispatch,
57 set: _set_dispatch
58 }
59
60
61 def _pickle_object(obj):
62 t = type(obj)
63 disp = _type_dispatch.get(t)
64 if disp is _identity_dispatch:
65 return obj
66
67 if disp is not None:
68 return disp(obj, _pickle_object)
69
70 if isinstance(obj, Exception):
71 return obj
72
73 getter = getattr(obj, '__getstate__', None)
74 if getter is not None:
75 state = getter()
76 else:
77 state = obj.__dict__
78
79 state = _dict_dispatch(state, _pickle_object)
80 state['__class__'] = obj.__class__.__name__
81 state['__module__'] = obj.__class__.__module__
82
83 return state
84
85
86 def _unpickle_object(state):
87 t = type(state)
88 disp = _type_dispatch.get(t)
89 if disp is _identity_dispatch:
90 return state
91
92 if (disp is not None and
93 (t != dict or '__module__' not in state)):
94 return disp(state, _unpickle_object)
95
96 if isinstance(state, Exception):
97 return state
98
99 mod_name = state['__module__']
100 mod = sys.modules[mod_name]
101 class_name = state['__class__']
102 class_def = getattr(mod, class_name)
103 obj = class_def.__new__(class_def)
104
105 del state['__class__']
106 del state['__module__']
107 attr_names = list(state.keys())
108 for name in attr_names:
109 state[name] = _unpickle_object(state[name])
110
111 obj.__dict__.update(state)
112
113 return obj
114