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