##// END OF EJS Templates
use brief dtype, rather than full object
MinRK -
Show More
@@ -1,232 +1,232 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 sys
20 20 from types import FunctionType
21 21
22 22 try:
23 23 import cPickle as pickle
24 24 except ImportError:
25 25 import pickle
26 26
27 27 try:
28 28 import numpy
29 29 except:
30 30 numpy = None
31 31
32 32 import codeutil
33 33 import py3compat
34 34 from importstring import import_item
35 35
36 36 if py3compat.PY3:
37 37 buffer = memoryview
38 38
39 39 #-------------------------------------------------------------------------------
40 40 # Classes
41 41 #-------------------------------------------------------------------------------
42 42
43 43
44 44 class CannedObject(object):
45 45 def __init__(self, obj, keys=[]):
46 46 self.keys = keys
47 47 self.obj = copy.copy(obj)
48 48 for key in keys:
49 49 setattr(self.obj, key, can(getattr(obj, key)))
50 50
51 51 self.buffers = []
52 52
53 53 def get_object(self, g=None):
54 54 if g is None:
55 55 g = {}
56 56 for key in self.keys:
57 57 setattr(self.obj, key, uncan(getattr(self.obj, key), g))
58 58 return self.obj
59 59
60 60
61 61 class Reference(CannedObject):
62 62 """object for wrapping a remote reference by name."""
63 63 def __init__(self, name):
64 64 if not isinstance(name, basestring):
65 65 raise TypeError("illegal name: %r"%name)
66 66 self.name = name
67 67 self.buffers = []
68 68
69 69 def __repr__(self):
70 70 return "<Reference: %r>"%self.name
71 71
72 72 def get_object(self, g=None):
73 73 if g is None:
74 74 g = {}
75 75
76 76 return eval(self.name, g)
77 77
78 78
79 79 class CannedFunction(CannedObject):
80 80
81 81 def __init__(self, f):
82 82 self._check_type(f)
83 83 self.code = f.func_code
84 84 self.defaults = f.func_defaults
85 85 self.module = f.__module__ or '__main__'
86 86 self.__name__ = f.__name__
87 87 self.buffers = []
88 88
89 89 def _check_type(self, obj):
90 90 assert isinstance(obj, FunctionType), "Not a function type"
91 91
92 92 def get_object(self, g=None):
93 93 # try to load function back into its module:
94 94 if not self.module.startswith('__'):
95 95 try:
96 96 __import__(self.module)
97 97 except ImportError:
98 98 pass
99 99 else:
100 100 g = sys.modules[self.module].__dict__
101 101
102 102 if g is None:
103 103 g = {}
104 104 newFunc = FunctionType(self.code, g, self.__name__, self.defaults)
105 105 return newFunc
106 106
107 107
108 108 class CannedArray(CannedObject):
109 109 def __init__(self, obj):
110 110 self.shape = obj.shape
111 self.dtype = obj.dtype
111 self.dtype = obj.dtype.descr if obj.dtype.fields else obj.dtype.str
112 112 if sum(obj.shape) == 0:
113 113 # just pickle it
114 114 self.buffers = [pickle.dumps(obj, -1)]
115 115 else:
116 116 # ensure contiguous
117 117 obj = numpy.ascontiguousarray(obj, dtype=None)
118 118 self.buffers = [buffer(obj)]
119 119
120 120 def get_object(self, g=None):
121 121 data = self.buffers[0]
122 122 if sum(self.shape) == 0:
123 123 # no shape, we just pickled it
124 124 return pickle.loads(data)
125 125 else:
126 126 return numpy.frombuffer(data, dtype=self.dtype).reshape(self.shape)
127 127
128 128
129 129 class CannedBytes(CannedObject):
130 130 wrap = bytes
131 131 def __init__(self, obj):
132 132 self.buffers = [obj]
133 133
134 134 def get_object(self, g=None):
135 135 data = self.buffers[0]
136 136 return self.wrap(data)
137 137
138 138 def CannedBuffer(CannedBytes):
139 139 wrap = buffer
140 140
141 141 #-------------------------------------------------------------------------------
142 142 # Functions
143 143 #-------------------------------------------------------------------------------
144 144
145 145
146 146 def can(obj):
147 147 """prepare an object for pickling"""
148 148 for cls,canner in can_map.iteritems():
149 149 if isinstance(cls, basestring):
150 150 try:
151 151 cls = import_item(cls)
152 152 except Exception:
153 153 # not importable
154 154 print "not importable: %r" % cls
155 155 continue
156 156 if isinstance(obj, cls):
157 157 return canner(obj)
158 158 return obj
159 159
160 160 def can_dict(obj):
161 161 """can the *values* of a dict"""
162 162 if isinstance(obj, dict):
163 163 newobj = {}
164 164 for k, v in obj.iteritems():
165 165 newobj[k] = can(v)
166 166 return newobj
167 167 else:
168 168 return obj
169 169
170 170 def can_sequence(obj):
171 171 """can the elements of a sequence"""
172 172 if isinstance(obj, (list, tuple)):
173 173 t = type(obj)
174 174 return t([can(i) for i in obj])
175 175 else:
176 176 return obj
177 177
178 178 def uncan(obj, g=None):
179 179 """invert canning"""
180 180 for cls,uncanner in uncan_map.iteritems():
181 181 if isinstance(cls, basestring):
182 182 try:
183 183 cls = import_item(cls)
184 184 except Exception:
185 185 # not importable
186 186 print "not importable: %r" % cls
187 187 continue
188 188 if isinstance(obj, cls):
189 189 return uncanner(obj, g)
190 190 return obj
191 191
192 192 def uncan_dict(obj, g=None):
193 193 if isinstance(obj, dict):
194 194 newobj = {}
195 195 for k, v in obj.iteritems():
196 196 newobj[k] = uncan(v,g)
197 197 return newobj
198 198 else:
199 199 return obj
200 200
201 201 def uncan_sequence(obj, g=None):
202 202 if isinstance(obj, (list, tuple)):
203 203 t = type(obj)
204 204 return t([uncan(i,g) for i in obj])
205 205 else:
206 206 return obj
207 207
208 208
209 209 #-------------------------------------------------------------------------------
210 210 # API dictionary
211 211 #-------------------------------------------------------------------------------
212 212
213 213 # These dicts can be extended for custom serialization of new objects
214 214
215 215 can_map = {
216 216 'IPython.parallel.dependent' : lambda obj: CannedObject(obj, keys=('f','df')),
217 217 'numpy.ndarray' : CannedArray,
218 218 FunctionType : CannedFunction,
219 219 bytes : CannedBytes,
220 220 buffer : CannedBuffer,
221 221 # dict : can_dict,
222 222 # list : can_sequence,
223 223 # tuple : can_sequence,
224 224 }
225 225
226 226 uncan_map = {
227 227 CannedObject : lambda obj, g: obj.get_object(g),
228 228 # dict : uncan_dict,
229 229 # list : uncan_sequence,
230 230 # tuple : uncan_sequence,
231 231 }
232 232
General Comments 0
You need to be logged in to leave comments. Login now