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