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