##// END OF EJS Templates
Don't catch ImportError when trying to unpack module functions...
MinRK -
Show More
@@ -1,281 +1,277 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:
28 try:
29 import numpy
29 import numpy
30 except:
30 except:
31 numpy = None
31 numpy = None
32
32
33 import codeutil
33 import codeutil
34 import py3compat
34 import py3compat
35 from importstring import import_item
35 from importstring import import_item
36
36
37 from IPython.config import Application
37 from IPython.config import Application
38
38
39 if py3compat.PY3:
39 if py3compat.PY3:
40 buffer = memoryview
40 buffer = memoryview
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=[]):
48 def __init__(self, obj, keys=[]):
49 self.keys = keys
49 self.keys = keys
50 self.obj = copy.copy(obj)
50 self.obj = copy.copy(obj)
51 for key in keys:
51 for key in keys:
52 setattr(self.obj, key, can(getattr(obj, key)))
52 setattr(self.obj, key, can(getattr(obj, key)))
53
53
54 self.buffers = []
54 self.buffers = []
55
55
56 def get_object(self, g=None):
56 def get_object(self, g=None):
57 if g is None:
57 if g is None:
58 g = {}
58 g = {}
59 for key in self.keys:
59 for key in self.keys:
60 setattr(self.obj, key, uncan(getattr(self.obj, key), g))
60 setattr(self.obj, key, uncan(getattr(self.obj, key), g))
61 return self.obj
61 return self.obj
62
62
63
63
64 class Reference(CannedObject):
64 class Reference(CannedObject):
65 """object for wrapping a remote reference by name."""
65 """object for wrapping a remote reference by name."""
66 def __init__(self, name):
66 def __init__(self, name):
67 if not isinstance(name, basestring):
67 if not isinstance(name, basestring):
68 raise TypeError("illegal name: %r"%name)
68 raise TypeError("illegal name: %r"%name)
69 self.name = name
69 self.name = name
70 self.buffers = []
70 self.buffers = []
71
71
72 def __repr__(self):
72 def __repr__(self):
73 return "<Reference: %r>"%self.name
73 return "<Reference: %r>"%self.name
74
74
75 def get_object(self, g=None):
75 def get_object(self, g=None):
76 if g is None:
76 if g is None:
77 g = {}
77 g = {}
78
78
79 return eval(self.name, g)
79 return eval(self.name, g)
80
80
81
81
82 class CannedFunction(CannedObject):
82 class CannedFunction(CannedObject):
83
83
84 def __init__(self, f):
84 def __init__(self, f):
85 self._check_type(f)
85 self._check_type(f)
86 self.code = f.func_code
86 self.code = f.func_code
87 if f.func_defaults:
87 if f.func_defaults:
88 self.defaults = [ can(fd) for fd in f.func_defaults ]
88 self.defaults = [ can(fd) for fd in f.func_defaults ]
89 else:
89 else:
90 self.defaults = None
90 self.defaults = None
91 self.module = f.__module__ or '__main__'
91 self.module = f.__module__ or '__main__'
92 self.__name__ = f.__name__
92 self.__name__ = f.__name__
93 self.buffers = []
93 self.buffers = []
94
94
95 def _check_type(self, obj):
95 def _check_type(self, obj):
96 assert isinstance(obj, FunctionType), "Not a function type"
96 assert isinstance(obj, FunctionType), "Not a function type"
97
97
98 def get_object(self, g=None):
98 def get_object(self, g=None):
99 # try to load function back into its module:
99 # try to load function back into its module:
100 if not self.module.startswith('__'):
100 if not self.module.startswith('__'):
101 try:
101 __import__(self.module)
102 __import__(self.module)
102 g = sys.modules[self.module].__dict__
103 except ImportError:
104 pass
105 else:
106 g = sys.modules[self.module].__dict__
107
103
108 if g is None:
104 if g is None:
109 g = {}
105 g = {}
110 if self.defaults:
106 if self.defaults:
111 defaults = tuple(uncan(cfd, g) for cfd in self.defaults)
107 defaults = tuple(uncan(cfd, g) for cfd in self.defaults)
112 else:
108 else:
113 defaults = None
109 defaults = None
114 newFunc = FunctionType(self.code, g, self.__name__, defaults)
110 newFunc = FunctionType(self.code, g, self.__name__, defaults)
115 return newFunc
111 return newFunc
116
112
117
113
118 class CannedArray(CannedObject):
114 class CannedArray(CannedObject):
119 def __init__(self, obj):
115 def __init__(self, obj):
120 self.shape = obj.shape
116 self.shape = obj.shape
121 self.dtype = obj.dtype.descr if obj.dtype.fields else obj.dtype.str
117 self.dtype = obj.dtype.descr if obj.dtype.fields else obj.dtype.str
122 if sum(obj.shape) == 0:
118 if sum(obj.shape) == 0:
123 # just pickle it
119 # just pickle it
124 self.buffers = [pickle.dumps(obj, -1)]
120 self.buffers = [pickle.dumps(obj, -1)]
125 else:
121 else:
126 # ensure contiguous
122 # ensure contiguous
127 obj = numpy.ascontiguousarray(obj, dtype=None)
123 obj = numpy.ascontiguousarray(obj, dtype=None)
128 self.buffers = [buffer(obj)]
124 self.buffers = [buffer(obj)]
129
125
130 def get_object(self, g=None):
126 def get_object(self, g=None):
131 data = self.buffers[0]
127 data = self.buffers[0]
132 if sum(self.shape) == 0:
128 if sum(self.shape) == 0:
133 # no shape, we just pickled it
129 # no shape, we just pickled it
134 return pickle.loads(data)
130 return pickle.loads(data)
135 else:
131 else:
136 return numpy.frombuffer(data, dtype=self.dtype).reshape(self.shape)
132 return numpy.frombuffer(data, dtype=self.dtype).reshape(self.shape)
137
133
138
134
139 class CannedBytes(CannedObject):
135 class CannedBytes(CannedObject):
140 wrap = bytes
136 wrap = bytes
141 def __init__(self, obj):
137 def __init__(self, obj):
142 self.buffers = [obj]
138 self.buffers = [obj]
143
139
144 def get_object(self, g=None):
140 def get_object(self, g=None):
145 data = self.buffers[0]
141 data = self.buffers[0]
146 return self.wrap(data)
142 return self.wrap(data)
147
143
148 def CannedBuffer(CannedBytes):
144 def CannedBuffer(CannedBytes):
149 wrap = buffer
145 wrap = buffer
150
146
151 #-------------------------------------------------------------------------------
147 #-------------------------------------------------------------------------------
152 # Functions
148 # Functions
153 #-------------------------------------------------------------------------------
149 #-------------------------------------------------------------------------------
154
150
155 def _logger():
151 def _logger():
156 """get the logger for the current Application
152 """get the logger for the current Application
157
153
158 the root logger will be used if no Application is running
154 the root logger will be used if no Application is running
159 """
155 """
160 if Application.initialized():
156 if Application.initialized():
161 logger = Application.instance().log
157 logger = Application.instance().log
162 else:
158 else:
163 logger = logging.getLogger()
159 logger = logging.getLogger()
164 if not logger.handlers:
160 if not logger.handlers:
165 logging.basicConfig()
161 logging.basicConfig()
166
162
167 return logger
163 return logger
168
164
169 def _import_mapping(mapping, original=None):
165 def _import_mapping(mapping, original=None):
170 """import any string-keys in a type mapping
166 """import any string-keys in a type mapping
171
167
172 """
168 """
173 log = _logger()
169 log = _logger()
174 log.debug("Importing canning map")
170 log.debug("Importing canning map")
175 for key,value in mapping.items():
171 for key,value in mapping.items():
176 if isinstance(key, basestring):
172 if isinstance(key, basestring):
177 try:
173 try:
178 cls = import_item(key)
174 cls = import_item(key)
179 except Exception:
175 except Exception:
180 if original and key not in original:
176 if original and key not in original:
181 # only message on user-added classes
177 # only message on user-added classes
182 log.error("cannning class not importable: %r", key, exc_info=True)
178 log.error("cannning class not importable: %r", key, exc_info=True)
183 mapping.pop(key)
179 mapping.pop(key)
184 else:
180 else:
185 mapping[cls] = mapping.pop(key)
181 mapping[cls] = mapping.pop(key)
186
182
187 def can(obj):
183 def can(obj):
188 """prepare an object for pickling"""
184 """prepare an object for pickling"""
189
185
190 import_needed = False
186 import_needed = False
191
187
192 for cls,canner in can_map.iteritems():
188 for cls,canner in can_map.iteritems():
193 if isinstance(cls, basestring):
189 if isinstance(cls, basestring):
194 import_needed = True
190 import_needed = True
195 break
191 break
196 elif isinstance(obj, cls):
192 elif isinstance(obj, cls):
197 return canner(obj)
193 return canner(obj)
198
194
199 if import_needed:
195 if import_needed:
200 # perform can_map imports, then try again
196 # perform can_map imports, then try again
201 # this will usually only happen once
197 # this will usually only happen once
202 _import_mapping(can_map, _original_can_map)
198 _import_mapping(can_map, _original_can_map)
203 return can(obj)
199 return can(obj)
204
200
205 return obj
201 return obj
206
202
207 def can_dict(obj):
203 def can_dict(obj):
208 """can the *values* of a dict"""
204 """can the *values* of a dict"""
209 if isinstance(obj, dict):
205 if isinstance(obj, dict):
210 newobj = {}
206 newobj = {}
211 for k, v in obj.iteritems():
207 for k, v in obj.iteritems():
212 newobj[k] = can(v)
208 newobj[k] = can(v)
213 return newobj
209 return newobj
214 else:
210 else:
215 return obj
211 return obj
216
212
217 def can_sequence(obj):
213 def can_sequence(obj):
218 """can the elements of a sequence"""
214 """can the elements of a sequence"""
219 if isinstance(obj, (list, tuple)):
215 if isinstance(obj, (list, tuple)):
220 t = type(obj)
216 t = type(obj)
221 return t([can(i) for i in obj])
217 return t([can(i) for i in obj])
222 else:
218 else:
223 return obj
219 return obj
224
220
225 def uncan(obj, g=None):
221 def uncan(obj, g=None):
226 """invert canning"""
222 """invert canning"""
227
223
228 import_needed = False
224 import_needed = False
229 for cls,uncanner in uncan_map.iteritems():
225 for cls,uncanner in uncan_map.iteritems():
230 if isinstance(cls, basestring):
226 if isinstance(cls, basestring):
231 import_needed = True
227 import_needed = True
232 break
228 break
233 elif isinstance(obj, cls):
229 elif isinstance(obj, cls):
234 return uncanner(obj, g)
230 return uncanner(obj, g)
235
231
236 if import_needed:
232 if import_needed:
237 # perform uncan_map imports, then try again
233 # perform uncan_map imports, then try again
238 # this will usually only happen once
234 # this will usually only happen once
239 _import_mapping(uncan_map, _original_uncan_map)
235 _import_mapping(uncan_map, _original_uncan_map)
240 return uncan(obj, g)
236 return uncan(obj, g)
241
237
242 return obj
238 return obj
243
239
244 def uncan_dict(obj, g=None):
240 def uncan_dict(obj, g=None):
245 if isinstance(obj, dict):
241 if isinstance(obj, dict):
246 newobj = {}
242 newobj = {}
247 for k, v in obj.iteritems():
243 for k, v in obj.iteritems():
248 newobj[k] = uncan(v,g)
244 newobj[k] = uncan(v,g)
249 return newobj
245 return newobj
250 else:
246 else:
251 return obj
247 return obj
252
248
253 def uncan_sequence(obj, g=None):
249 def uncan_sequence(obj, g=None):
254 if isinstance(obj, (list, tuple)):
250 if isinstance(obj, (list, tuple)):
255 t = type(obj)
251 t = type(obj)
256 return t([uncan(i,g) for i in obj])
252 return t([uncan(i,g) for i in obj])
257 else:
253 else:
258 return obj
254 return obj
259
255
260
256
261 #-------------------------------------------------------------------------------
257 #-------------------------------------------------------------------------------
262 # API dictionaries
258 # API dictionaries
263 #-------------------------------------------------------------------------------
259 #-------------------------------------------------------------------------------
264
260
265 # These dicts can be extended for custom serialization of new objects
261 # These dicts can be extended for custom serialization of new objects
266
262
267 can_map = {
263 can_map = {
268 'IPython.parallel.dependent' : lambda obj: CannedObject(obj, keys=('f','df')),
264 'IPython.parallel.dependent' : lambda obj: CannedObject(obj, keys=('f','df')),
269 'numpy.ndarray' : CannedArray,
265 'numpy.ndarray' : CannedArray,
270 FunctionType : CannedFunction,
266 FunctionType : CannedFunction,
271 bytes : CannedBytes,
267 bytes : CannedBytes,
272 buffer : CannedBuffer,
268 buffer : CannedBuffer,
273 }
269 }
274
270
275 uncan_map = {
271 uncan_map = {
276 CannedObject : lambda obj, g: obj.get_object(g),
272 CannedObject : lambda obj, g: obj.get_object(g),
277 }
273 }
278
274
279 # for use in _import_mapping:
275 # for use in _import_mapping:
280 _original_can_map = can_map.copy()
276 _original_can_map = can_map.copy()
281 _original_uncan_map = uncan_map.copy()
277 _original_uncan_map = uncan_map.copy()
General Comments 0
You need to be logged in to leave comments. Login now