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