Mercurial > piecrust2

comparison piecrust/fastpickle.py @ 461:b015e38d4ee1

Find changesets by keywords (author, files, the commit message), revision number or hash, or revset expression.
internal: Handle data serialization more under the hood. Implement a new queue class that handles pickling of objects, add option (currently disabled) to also do zlib compression before sending bytes.
author Ludovic Chabant <ludovic@chabant.com>
date Sat, 11 Jul 2015 17:51:56 -0700
parents 2ef04e16f0b9
children 9e5393fcfab2
comparison
equal deleted inserted replaced
460:55fc8918cb75 461:b015e38d4ee1
1 import sys 1 import sys
2 import json
2 import datetime 3 import datetime
3 import collections 4 import collections
4 5
5 6
6 def pickle(obj): 7 def pickle(obj):
7 return _pickle_object(obj) 8 data = _pickle_object(obj)
9 data = json.dumps(data, indent=None, separators=(',', ':'))
10 return data.encode('utf8')
8 11
9 12
10 def unpickle(data): 13 def unpickle(data):
14 data = data.decode('utf8')
15 data = json.loads(data)
11 return _unpickle_object(data) 16 return _unpickle_object(data)
12 17
13 18
14 _PICKLING = 0 19 _PICKLING = 0
15 _UNPICKLING = 1 20 _UNPICKLING = 1
16 21
17 22 _identity_dispatch = object()
18 def _tuple_dispatch(obj, func, op): 23
19 res = [None] * len(obj) 24
20 for i, c in enumerate(obj): 25 def _tuple_convert(obj, func, op):
21 res[i] = func(c) 26 if op == _PICKLING:
22 return tuple(res) 27 return ['__type__:tuple'] + [func(c) for c in obj]
23 28 elif op == _UNPICKLING:
24 29 return tuple([func(c) for c in obj[1:]])
25 def _list_dispatch(obj, func, op): 30
26 res = [None] * len(obj) 31
27 for i, c in enumerate(obj): 32 def _list_convert(obj, func, op):
28 res[i] = func(c) 33 return [func(c) for c in obj]
29 return res 34
30 35
31 36 def _dict_convert(obj, func, op):
32 def _dict_dispatch(obj, func, op):
33 res = {} 37 res = {}
34 for k, v in obj.items(): 38 for k, v in obj.items():
35 res[k] = func(v) 39 res[k] = func(v)
36 return res 40 return res
37 41
38 42
39 def _set_dispatch(obj, func, op): 43 def _ordered_dict_convert(obj, func, op):
40 res = set() 44 if op == _PICKLING:
41 for v in obj: 45 res = {'__type__': 'OrderedDict'}
42 res.add(func(v)) 46 for k, v in obj.items():
43 return res 47 res[k] = func(v)
44 48 return res
45 49 elif op == _UNPICKLING:
46 def _date_convert(obj, op): 50 res = collections.OrderedDict()
51 for k, v in obj.items():
52 res[k] = func(v)
53 return res
54
55
56 def _set_convert(obj, func, op):
57 if op == _PICKLING:
58 return ['__type__:set'] + [func(c) for c in obj]
59 elif op == _UNPICKLING:
60 return set([func(c) for c in obj[1:]])
61
62
63 def _date_convert(obj, func, op):
47 if op == _PICKLING: 64 if op == _PICKLING:
48 return {'__class__': 'date', 65 return {'__class__': 'date',
49 'year': obj.year, 66 'year': obj.year,
50 'month': obj.month, 67 'month': obj.month,
51 'day': obj.day} 68 'day': obj.day}
52 elif op == _UNPICKLING: 69 elif op == _UNPICKLING:
53 return datetime.date( 70 return datetime.date(
54 obj['year'], obj['month'], obj['day']) 71 obj['year'], obj['month'], obj['day'])
55 72
56 73
57 def _datetime_convert(obj, op): 74 def _datetime_convert(obj, func, op):
58 if op == _PICKLING: 75 if op == _PICKLING:
59 return {'__class__': 'datetime', 76 return {'__class__': 'datetime',
60 'year': obj.year, 77 'year': obj.year,
61 'month': obj.month, 78 'month': obj.month,
62 'day': obj.day, 79 'day': obj.day,
68 return datetime.datetime( 85 return datetime.datetime(
69 obj['year'], obj['month'], obj['day'], 86 obj['year'], obj['month'], obj['day'],
70 obj['hour'], obj['minute'], obj['second'], obj['microsecond']) 87 obj['hour'], obj['minute'], obj['second'], obj['microsecond'])
71 88
72 89
73 def _time_convert(obj, op): 90 def _time_convert(obj, func, op):
74 if op == _PICKLING: 91 if op == _PICKLING:
75 return {'__class__': 'time', 92 return {'__class__': 'time',
76 'hour': obj.hour, 93 'hour': obj.hour,
77 'minute': obj.minute, 94 'minute': obj.minute,
78 'second': obj.second, 95 'second': obj.second,
80 elif op == _UNPICKLING: 97 elif op == _UNPICKLING:
81 return datetime.time( 98 return datetime.time(
82 obj['hour'], obj['minute'], obj['second'], obj['microsecond']) 99 obj['hour'], obj['minute'], obj['second'], obj['microsecond'])
83 100
84 101
85 _identity_dispatch = object() 102 _type_convert = {
86
87 _type_dispatch = {
88 type(None): _identity_dispatch, 103 type(None): _identity_dispatch,
89 bool: _identity_dispatch, 104 bool: _identity_dispatch,
90 int: _identity_dispatch, 105 int: _identity_dispatch,
91 float: _identity_dispatch, 106 float: _identity_dispatch,
92 str: _identity_dispatch, 107 str: _identity_dispatch,
93 tuple: _tuple_dispatch,
94 list: _list_dispatch,
95 dict: _dict_dispatch,
96 collections.OrderedDict: _dict_dispatch,
97 set: _set_dispatch
98 }
99
100
101 _type_convert = {
102 datetime.date: _date_convert, 108 datetime.date: _date_convert,
103 datetime.datetime: _datetime_convert, 109 datetime.datetime: _datetime_convert,
104 datetime.time: _time_convert 110 datetime.time: _time_convert,
111 tuple: _tuple_convert,
112 list: _list_convert,
113 dict: _dict_convert,
114 set: _set_convert,
115 collections.OrderedDict: _ordered_dict_convert,
105 } 116 }
106 117
107 118
108 _type_unconvert = { 119 _type_unconvert = {
120 type(None): _identity_dispatch,
121 bool: _identity_dispatch,
122 int: _identity_dispatch,
123 float: _identity_dispatch,
124 str: _identity_dispatch,
109 'date': _date_convert, 125 'date': _date_convert,
110 'datetime': _datetime_convert, 126 'datetime': _datetime_convert,
111 'time': _time_convert 127 'time': _time_convert,
128 }
129
130
131 _collection_unconvert = {
132 '__type__:tuple': _tuple_convert,
133 '__type__:set': _set_convert,
134 }
135
136
137 _mapping_unconvert = {
138 'OrderedDict': _ordered_dict_convert
112 } 139 }
113 140
114 141
115 def _pickle_object(obj): 142 def _pickle_object(obj):
116 t = type(obj) 143 t = type(obj)
117 disp = _type_dispatch.get(t) 144 conv = _type_convert.get(t)
118 if disp is _identity_dispatch: 145
146 # Object doesn't need conversion?
147 if conv is _identity_dispatch:
119 return obj 148 return obj
120 149
121 if disp is not None: 150 # Object has special conversion?
122 return disp(obj, _pickle_object, _PICKLING)
123
124 conv = _type_convert.get(t)
125 if conv is not None: 151 if conv is not None:
126 return conv(obj, _PICKLING) 152 return conv(obj, _pickle_object, _PICKLING)
127 153
154 # Use instance dictionary, or a custom state.
128 getter = getattr(obj, '__getstate__', None) 155 getter = getattr(obj, '__getstate__', None)
129 if getter is not None: 156 if getter is not None:
130 state = getter() 157 state = getter()
131 else: 158 else:
132 state = obj.__dict__ 159 state = obj.__dict__
133 160
134 state = _dict_dispatch(state, _pickle_object, _PICKLING) 161 state = _dict_convert(state, _pickle_object, _PICKLING)
135 state['__class__'] = obj.__class__.__name__ 162 state['__class__'] = obj.__class__.__name__
136 state['__module__'] = obj.__class__.__module__ 163 state['__module__'] = obj.__class__.__module__
137 164
138 return state 165 return state
139 166
140 167
141 def _unpickle_object(state): 168 def _unpickle_object(state):
142 t = type(state) 169 t = type(state)
143 disp = _type_dispatch.get(t) 170 conv = _type_unconvert.get(t)
144 if disp is _identity_dispatch: 171
172 # Object doesn't need conversion?
173 if conv is _identity_dispatch:
145 return state 174 return state
146 175
147 if (disp is not None and 176 # Try collection or mapping conversion.
148 (t != dict or '__class__' not in state)): 177 if t is list:
149 return disp(state, _unpickle_object, _UNPICKLING) 178 try:
150 179 col_type = state[0]
151 class_name = state['__class__'] 180 if not isinstance(col_type, str):
181 col_type = None
182 except IndexError:
183 col_type = None
184 if col_type is not None:
185 conv = _collection_unconvert.get(col_type)
186 if conv is not None:
187 return conv(state, _unpickle_object, _UNPICKLING)
188 return _list_convert(state, _unpickle_object, _UNPICKLING)
189
190 assert t is dict
191
192 # Custom mapping type?
193 map_type = state.get('__type__')
194 if map_type:
195 conv = _mapping_unconvert.get(map_type)
196 return conv(state, _unpickle_object, _UNPICKLING)
197
198 # Class instance or other custom type.
199 class_name = state.get('__class__')
200 if class_name is None:
201 return _dict_convert(state, _unpickle_object, _UNPICKLING)
202
152 conv = _type_unconvert.get(class_name) 203 conv = _type_unconvert.get(class_name)
153 if conv is not None: 204 if conv is not None:
154 return conv(state, _UNPICKLING) 205 return conv(state, _unpickle_object, _UNPICKLING)
155 206
156 mod_name = state['__module__'] 207 mod_name = state['__module__']
157 mod = sys.modules[mod_name] 208 mod = sys.modules[mod_name]
158 class_def = getattr(mod, class_name) 209 class_def = getattr(mod, class_name)
159 obj = class_def.__new__(class_def) 210 obj = class_def.__new__(class_def)