##// END OF EJS Templates
Fix for dictionary iteration in pickleutil
Thomas Kluyver -
Show More
@@ -1,352 +1,352
1 1 # encoding: utf-8
2 2
3 3 """Pickle related utilities. Perhaps this should be called 'can'."""
4 4
5 5 __docformat__ = "restructuredtext en"
6 6
7 7 #-------------------------------------------------------------------------------
8 8 # Copyright (C) 2008-2011 The IPython Development Team
9 9 #
10 10 # Distributed under the terms of the BSD License. The full license is in
11 11 # the file COPYING, distributed as part of this software.
12 12 #-------------------------------------------------------------------------------
13 13
14 14 #-------------------------------------------------------------------------------
15 15 # Imports
16 16 #-------------------------------------------------------------------------------
17 17
18 18 import copy
19 19 import logging
20 20 import sys
21 21 from types import FunctionType
22 22
23 23 try:
24 24 import cPickle as pickle
25 25 except ImportError:
26 26 import pickle
27 27
28 28 from . import codeutil # This registers a hook when it's imported
29 29 from . import py3compat
30 30 from .importstring import import_item
31 31 from .py3compat import string_types, iteritems
32 32
33 33 from IPython.config import Application
34 34
35 35 if py3compat.PY3:
36 36 buffer = memoryview
37 37 class_type = type
38 38 else:
39 39 from types import ClassType
40 40 class_type = (type, ClassType)
41 41
42 42 #-------------------------------------------------------------------------------
43 43 # Classes
44 44 #-------------------------------------------------------------------------------
45 45
46 46
47 47 class CannedObject(object):
48 48 def __init__(self, obj, keys=[], hook=None):
49 49 """can an object for safe pickling
50 50
51 51 Parameters
52 52 ==========
53 53
54 54 obj:
55 55 The object to be canned
56 56 keys: list (optional)
57 57 list of attribute names that will be explicitly canned / uncanned
58 58 hook: callable (optional)
59 59 An optional extra callable,
60 60 which can do additional processing of the uncanned object.
61 61
62 62 large data may be offloaded into the buffers list,
63 63 used for zero-copy transfers.
64 64 """
65 65 self.keys = keys
66 66 self.obj = copy.copy(obj)
67 67 self.hook = can(hook)
68 68 for key in keys:
69 69 setattr(self.obj, key, can(getattr(obj, key)))
70 70
71 71 self.buffers = []
72 72
73 73 def get_object(self, g=None):
74 74 if g is None:
75 75 g = {}
76 76 obj = self.obj
77 77 for key in self.keys:
78 78 setattr(obj, key, uncan(getattr(obj, key), g))
79 79
80 80 if self.hook:
81 81 self.hook = uncan(self.hook, g)
82 82 self.hook(obj, g)
83 83 return self.obj
84 84
85 85
86 86 class Reference(CannedObject):
87 87 """object for wrapping a remote reference by name."""
88 88 def __init__(self, name):
89 89 if not isinstance(name, string_types):
90 90 raise TypeError("illegal name: %r"%name)
91 91 self.name = name
92 92 self.buffers = []
93 93
94 94 def __repr__(self):
95 95 return "<Reference: %r>"%self.name
96 96
97 97 def get_object(self, g=None):
98 98 if g is None:
99 99 g = {}
100 100
101 101 return eval(self.name, g)
102 102
103 103
104 104 class CannedFunction(CannedObject):
105 105
106 106 def __init__(self, f):
107 107 self._check_type(f)
108 108 self.code = f.__code__
109 109 if f.__defaults__:
110 110 self.defaults = [ can(fd) for fd in f.__defaults__ ]
111 111 else:
112 112 self.defaults = None
113 113 self.module = f.__module__ or '__main__'
114 114 self.__name__ = f.__name__
115 115 self.buffers = []
116 116
117 117 def _check_type(self, obj):
118 118 assert isinstance(obj, FunctionType), "Not a function type"
119 119
120 120 def get_object(self, g=None):
121 121 # try to load function back into its module:
122 122 if not self.module.startswith('__'):
123 123 __import__(self.module)
124 124 g = sys.modules[self.module].__dict__
125 125
126 126 if g is None:
127 127 g = {}
128 128 if self.defaults:
129 129 defaults = tuple(uncan(cfd, g) for cfd in self.defaults)
130 130 else:
131 131 defaults = None
132 132 newFunc = FunctionType(self.code, g, self.__name__, defaults)
133 133 return newFunc
134 134
135 135 class CannedClass(CannedObject):
136 136
137 137 def __init__(self, cls):
138 138 self._check_type(cls)
139 139 self.name = cls.__name__
140 140 self.old_style = not isinstance(cls, type)
141 141 self._canned_dict = {}
142 142 for k,v in cls.__dict__.items():
143 143 if k not in ('__weakref__', '__dict__'):
144 144 self._canned_dict[k] = can(v)
145 145 if self.old_style:
146 146 mro = []
147 147 else:
148 148 mro = cls.mro()
149 149
150 150 self.parents = [ can(c) for c in mro[1:] ]
151 151 self.buffers = []
152 152
153 153 def _check_type(self, obj):
154 154 assert isinstance(obj, class_type), "Not a class type"
155 155
156 156 def get_object(self, g=None):
157 157 parents = tuple(uncan(p, g) for p in self.parents)
158 158 return type(self.name, parents, uncan_dict(self._canned_dict, g=g))
159 159
160 160 class CannedArray(CannedObject):
161 161 def __init__(self, obj):
162 162 from numpy import ascontiguousarray
163 163 self.shape = obj.shape
164 164 self.dtype = obj.dtype.descr if obj.dtype.fields else obj.dtype.str
165 165 if sum(obj.shape) == 0:
166 166 # just pickle it
167 167 self.buffers = [pickle.dumps(obj, -1)]
168 168 else:
169 169 # ensure contiguous
170 170 obj = ascontiguousarray(obj, dtype=None)
171 171 self.buffers = [buffer(obj)]
172 172
173 173 def get_object(self, g=None):
174 174 from numpy import frombuffer
175 175 data = self.buffers[0]
176 176 if sum(self.shape) == 0:
177 177 # no shape, we just pickled it
178 178 return pickle.loads(data)
179 179 else:
180 180 return frombuffer(data, dtype=self.dtype).reshape(self.shape)
181 181
182 182
183 183 class CannedBytes(CannedObject):
184 184 wrap = bytes
185 185 def __init__(self, obj):
186 186 self.buffers = [obj]
187 187
188 188 def get_object(self, g=None):
189 189 data = self.buffers[0]
190 190 return self.wrap(data)
191 191
192 192 def CannedBuffer(CannedBytes):
193 193 wrap = buffer
194 194
195 195 #-------------------------------------------------------------------------------
196 196 # Functions
197 197 #-------------------------------------------------------------------------------
198 198
199 199 def _logger():
200 200 """get the logger for the current Application
201 201
202 202 the root logger will be used if no Application is running
203 203 """
204 204 if Application.initialized():
205 205 logger = Application.instance().log
206 206 else:
207 207 logger = logging.getLogger()
208 208 if not logger.handlers:
209 209 logging.basicConfig()
210 210
211 211 return logger
212 212
213 213 def _import_mapping(mapping, original=None):
214 214 """import any string-keys in a type mapping
215 215
216 216 """
217 217 log = _logger()
218 218 log.debug("Importing canning map")
219 for key,value in mapping.items():
219 for key,value in list(mapping.items()):
220 220 if isinstance(key, string_types):
221 221 try:
222 222 cls = import_item(key)
223 223 except Exception:
224 224 if original and key not in original:
225 225 # only message on user-added classes
226 226 log.error("canning class not importable: %r", key, exc_info=True)
227 227 mapping.pop(key)
228 228 else:
229 229 mapping[cls] = mapping.pop(key)
230 230
231 231 def istype(obj, check):
232 232 """like isinstance(obj, check), but strict
233 233
234 234 This won't catch subclasses.
235 235 """
236 236 if isinstance(check, tuple):
237 237 for cls in check:
238 238 if type(obj) is cls:
239 239 return True
240 240 return False
241 241 else:
242 242 return type(obj) is check
243 243
244 244 def can(obj):
245 245 """prepare an object for pickling"""
246 246
247 247 import_needed = False
248 248
249 249 for cls,canner in iteritems(can_map):
250 250 if isinstance(cls, string_types):
251 251 import_needed = True
252 252 break
253 253 elif istype(obj, cls):
254 254 return canner(obj)
255 255
256 256 if import_needed:
257 257 # perform can_map imports, then try again
258 258 # this will usually only happen once
259 259 _import_mapping(can_map, _original_can_map)
260 260 return can(obj)
261 261
262 262 return obj
263 263
264 264 def can_class(obj):
265 265 if isinstance(obj, class_type) and obj.__module__ == '__main__':
266 266 return CannedClass(obj)
267 267 else:
268 268 return obj
269 269
270 270 def can_dict(obj):
271 271 """can the *values* of a dict"""
272 272 if istype(obj, dict):
273 273 newobj = {}
274 274 for k, v in iteritems(obj):
275 275 newobj[k] = can(v)
276 276 return newobj
277 277 else:
278 278 return obj
279 279
280 280 sequence_types = (list, tuple, set)
281 281
282 282 def can_sequence(obj):
283 283 """can the elements of a sequence"""
284 284 if istype(obj, sequence_types):
285 285 t = type(obj)
286 286 return t([can(i) for i in obj])
287 287 else:
288 288 return obj
289 289
290 290 def uncan(obj, g=None):
291 291 """invert canning"""
292 292
293 293 import_needed = False
294 294 for cls,uncanner in iteritems(uncan_map):
295 295 if isinstance(cls, string_types):
296 296 import_needed = True
297 297 break
298 298 elif isinstance(obj, cls):
299 299 return uncanner(obj, g)
300 300
301 301 if import_needed:
302 302 # perform uncan_map imports, then try again
303 303 # this will usually only happen once
304 304 _import_mapping(uncan_map, _original_uncan_map)
305 305 return uncan(obj, g)
306 306
307 307 return obj
308 308
309 309 def uncan_dict(obj, g=None):
310 310 if istype(obj, dict):
311 311 newobj = {}
312 312 for k, v in iteritems(obj):
313 313 newobj[k] = uncan(v,g)
314 314 return newobj
315 315 else:
316 316 return obj
317 317
318 318 def uncan_sequence(obj, g=None):
319 319 if istype(obj, sequence_types):
320 320 t = type(obj)
321 321 return t([uncan(i,g) for i in obj])
322 322 else:
323 323 return obj
324 324
325 325 def _uncan_dependent_hook(dep, g=None):
326 326 dep.check_dependency()
327 327
328 328 def can_dependent(obj):
329 329 return CannedObject(obj, keys=('f', 'df'), hook=_uncan_dependent_hook)
330 330
331 331 #-------------------------------------------------------------------------------
332 332 # API dictionaries
333 333 #-------------------------------------------------------------------------------
334 334
335 335 # These dicts can be extended for custom serialization of new objects
336 336
337 337 can_map = {
338 338 'IPython.parallel.dependent' : can_dependent,
339 339 'numpy.ndarray' : CannedArray,
340 340 FunctionType : CannedFunction,
341 341 bytes : CannedBytes,
342 342 buffer : CannedBuffer,
343 343 class_type : can_class,
344 344 }
345 345
346 346 uncan_map = {
347 347 CannedObject : lambda obj, g: obj.get_object(g),
348 348 }
349 349
350 350 # for use in _import_mapping:
351 351 _original_can_map = can_map.copy()
352 352 _original_uncan_map = uncan_map.copy()
General Comments 0
You need to be logged in to leave comments. Login now