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