Show More
@@ -1,185 +1,181 b'' | |||||
1 | """serialization utilities for apply messages""" |
|
1 | """serialization utilities for apply messages""" | |
2 |
|
2 | |||
3 | # Copyright (c) IPython Development Team. |
|
3 | # Copyright (c) IPython Development Team. | |
4 | # Distributed under the terms of the Modified BSD License. |
|
4 | # Distributed under the terms of the Modified BSD License. | |
5 |
|
5 | |||
6 | try: |
|
6 | try: | |
7 | import cPickle |
|
7 | import cPickle | |
8 | pickle = cPickle |
|
8 | pickle = cPickle | |
9 | except: |
|
9 | except: | |
10 | cPickle = None |
|
10 | cPickle = None | |
11 | import pickle |
|
11 | import pickle | |
12 |
|
12 | |||
13 | # IPython imports |
|
13 | # IPython imports | |
14 | from IPython.utils import py3compat |
|
14 | from IPython.utils.py3compat import PY3, buffer_to_bytes_py2 | |
15 | from IPython.utils.data import flatten |
|
15 | from IPython.utils.data import flatten | |
16 | from IPython.utils.pickleutil import ( |
|
16 | from IPython.utils.pickleutil import ( | |
17 | can, uncan, can_sequence, uncan_sequence, CannedObject, |
|
17 | can, uncan, can_sequence, uncan_sequence, CannedObject, | |
18 | istype, sequence_types, PICKLE_PROTOCOL, |
|
18 | istype, sequence_types, PICKLE_PROTOCOL, | |
19 | ) |
|
19 | ) | |
20 |
|
20 | |||
21 | if py3compat.PY3: |
|
21 | if PY3: | |
22 | buffer = memoryview |
|
22 | buffer = memoryview | |
23 |
|
23 | |||
24 | #----------------------------------------------------------------------------- |
|
24 | #----------------------------------------------------------------------------- | |
25 | # Serialization Functions |
|
25 | # Serialization Functions | |
26 | #----------------------------------------------------------------------------- |
|
26 | #----------------------------------------------------------------------------- | |
27 |
|
27 | |||
28 | # default values for the thresholds: |
|
28 | # default values for the thresholds: | |
29 | MAX_ITEMS = 64 |
|
29 | MAX_ITEMS = 64 | |
30 | MAX_BYTES = 1024 |
|
30 | MAX_BYTES = 1024 | |
31 |
|
31 | |||
32 | def _extract_buffers(obj, threshold=MAX_BYTES): |
|
32 | def _extract_buffers(obj, threshold=MAX_BYTES): | |
33 | """extract buffers larger than a certain threshold""" |
|
33 | """extract buffers larger than a certain threshold""" | |
34 | buffers = [] |
|
34 | buffers = [] | |
35 | if isinstance(obj, CannedObject) and obj.buffers: |
|
35 | if isinstance(obj, CannedObject) and obj.buffers: | |
36 | for i,buf in enumerate(obj.buffers): |
|
36 | for i,buf in enumerate(obj.buffers): | |
37 | if len(buf) > threshold: |
|
37 | if len(buf) > threshold: | |
38 | # buffer larger than threshold, prevent pickling |
|
38 | # buffer larger than threshold, prevent pickling | |
39 | obj.buffers[i] = None |
|
39 | obj.buffers[i] = None | |
40 | buffers.append(buf) |
|
40 | buffers.append(buf) | |
41 | elif isinstance(buf, buffer): |
|
41 | elif isinstance(buf, buffer): | |
42 | # buffer too small for separate send, coerce to bytes |
|
42 | # buffer too small for separate send, coerce to bytes | |
43 | # because pickling buffer objects just results in broken pointers |
|
43 | # because pickling buffer objects just results in broken pointers | |
44 | obj.buffers[i] = bytes(buf) |
|
44 | obj.buffers[i] = bytes(buf) | |
45 | return buffers |
|
45 | return buffers | |
46 |
|
46 | |||
47 | def _restore_buffers(obj, buffers): |
|
47 | def _restore_buffers(obj, buffers): | |
48 | """restore buffers extracted by """ |
|
48 | """restore buffers extracted by """ | |
49 | if isinstance(obj, CannedObject) and obj.buffers: |
|
49 | if isinstance(obj, CannedObject) and obj.buffers: | |
50 | for i,buf in enumerate(obj.buffers): |
|
50 | for i,buf in enumerate(obj.buffers): | |
51 | if buf is None: |
|
51 | if buf is None: | |
52 | obj.buffers[i] = buffers.pop(0) |
|
52 | obj.buffers[i] = buffers.pop(0) | |
53 |
|
53 | |||
54 | def serialize_object(obj, buffer_threshold=MAX_BYTES, item_threshold=MAX_ITEMS): |
|
54 | def serialize_object(obj, buffer_threshold=MAX_BYTES, item_threshold=MAX_ITEMS): | |
55 | """Serialize an object into a list of sendable buffers. |
|
55 | """Serialize an object into a list of sendable buffers. | |
56 |
|
56 | |||
57 | Parameters |
|
57 | Parameters | |
58 | ---------- |
|
58 | ---------- | |
59 |
|
59 | |||
60 | obj : object |
|
60 | obj : object | |
61 | The object to be serialized |
|
61 | The object to be serialized | |
62 | buffer_threshold : int |
|
62 | buffer_threshold : int | |
63 | The threshold (in bytes) for pulling out data buffers |
|
63 | The threshold (in bytes) for pulling out data buffers | |
64 | to avoid pickling them. |
|
64 | to avoid pickling them. | |
65 | item_threshold : int |
|
65 | item_threshold : int | |
66 | The maximum number of items over which canning will iterate. |
|
66 | The maximum number of items over which canning will iterate. | |
67 | Containers (lists, dicts) larger than this will be pickled without |
|
67 | Containers (lists, dicts) larger than this will be pickled without | |
68 | introspection. |
|
68 | introspection. | |
69 |
|
69 | |||
70 | Returns |
|
70 | Returns | |
71 | ------- |
|
71 | ------- | |
72 | [bufs] : list of buffers representing the serialized object. |
|
72 | [bufs] : list of buffers representing the serialized object. | |
73 | """ |
|
73 | """ | |
74 | buffers = [] |
|
74 | buffers = [] | |
75 | if istype(obj, sequence_types) and len(obj) < item_threshold: |
|
75 | if istype(obj, sequence_types) and len(obj) < item_threshold: | |
76 | cobj = can_sequence(obj) |
|
76 | cobj = can_sequence(obj) | |
77 | for c in cobj: |
|
77 | for c in cobj: | |
78 | buffers.extend(_extract_buffers(c, buffer_threshold)) |
|
78 | buffers.extend(_extract_buffers(c, buffer_threshold)) | |
79 | elif istype(obj, dict) and len(obj) < item_threshold: |
|
79 | elif istype(obj, dict) and len(obj) < item_threshold: | |
80 | cobj = {} |
|
80 | cobj = {} | |
81 | for k in sorted(obj): |
|
81 | for k in sorted(obj): | |
82 | c = can(obj[k]) |
|
82 | c = can(obj[k]) | |
83 | buffers.extend(_extract_buffers(c, buffer_threshold)) |
|
83 | buffers.extend(_extract_buffers(c, buffer_threshold)) | |
84 | cobj[k] = c |
|
84 | cobj[k] = c | |
85 | else: |
|
85 | else: | |
86 | cobj = can(obj) |
|
86 | cobj = can(obj) | |
87 | buffers.extend(_extract_buffers(cobj, buffer_threshold)) |
|
87 | buffers.extend(_extract_buffers(cobj, buffer_threshold)) | |
88 |
|
88 | |||
89 | buffers.insert(0, pickle.dumps(cobj, PICKLE_PROTOCOL)) |
|
89 | buffers.insert(0, pickle.dumps(cobj, PICKLE_PROTOCOL)) | |
90 | return buffers |
|
90 | return buffers | |
91 |
|
91 | |||
92 | def deserialize_object(buffers, g=None): |
|
92 | def deserialize_object(buffers, g=None): | |
93 | """reconstruct an object serialized by serialize_object from data buffers. |
|
93 | """reconstruct an object serialized by serialize_object from data buffers. | |
94 |
|
94 | |||
95 | Parameters |
|
95 | Parameters | |
96 | ---------- |
|
96 | ---------- | |
97 |
|
97 | |||
98 | bufs : list of buffers/bytes |
|
98 | bufs : list of buffers/bytes | |
99 |
|
99 | |||
100 | g : globals to be used when uncanning |
|
100 | g : globals to be used when uncanning | |
101 |
|
101 | |||
102 | Returns |
|
102 | Returns | |
103 | ------- |
|
103 | ------- | |
104 |
|
104 | |||
105 | (newobj, bufs) : unpacked object, and the list of remaining unused buffers. |
|
105 | (newobj, bufs) : unpacked object, and the list of remaining unused buffers. | |
106 | """ |
|
106 | """ | |
107 | bufs = list(buffers) |
|
107 | bufs = list(buffers) | |
108 | pobj = bufs.pop(0) |
|
108 | pobj = buffer_to_bytes_py2(bufs.pop(0)) | |
109 | if not isinstance(pobj, bytes): |
|
|||
110 | # a zmq message |
|
|||
111 | pobj = bytes(pobj) |
|
|||
112 | canned = pickle.loads(pobj) |
|
109 | canned = pickle.loads(pobj) | |
113 | if istype(canned, sequence_types) and len(canned) < MAX_ITEMS: |
|
110 | if istype(canned, sequence_types) and len(canned) < MAX_ITEMS: | |
114 | for c in canned: |
|
111 | for c in canned: | |
115 | _restore_buffers(c, bufs) |
|
112 | _restore_buffers(c, bufs) | |
116 | newobj = uncan_sequence(canned, g) |
|
113 | newobj = uncan_sequence(canned, g) | |
117 | elif istype(canned, dict) and len(canned) < MAX_ITEMS: |
|
114 | elif istype(canned, dict) and len(canned) < MAX_ITEMS: | |
118 | newobj = {} |
|
115 | newobj = {} | |
119 | for k in sorted(canned): |
|
116 | for k in sorted(canned): | |
120 | c = canned[k] |
|
117 | c = canned[k] | |
121 | _restore_buffers(c, bufs) |
|
118 | _restore_buffers(c, bufs) | |
122 | newobj[k] = uncan(c, g) |
|
119 | newobj[k] = uncan(c, g) | |
123 | else: |
|
120 | else: | |
124 | _restore_buffers(canned, bufs) |
|
121 | _restore_buffers(canned, bufs) | |
125 | newobj = uncan(canned, g) |
|
122 | newobj = uncan(canned, g) | |
126 |
|
123 | |||
127 | return newobj, bufs |
|
124 | return newobj, bufs | |
128 |
|
125 | |||
129 | def pack_apply_message(f, args, kwargs, buffer_threshold=MAX_BYTES, item_threshold=MAX_ITEMS): |
|
126 | def pack_apply_message(f, args, kwargs, buffer_threshold=MAX_BYTES, item_threshold=MAX_ITEMS): | |
130 | """pack up a function, args, and kwargs to be sent over the wire |
|
127 | """pack up a function, args, and kwargs to be sent over the wire | |
131 |
|
128 | |||
132 | Each element of args/kwargs will be canned for special treatment, |
|
129 | Each element of args/kwargs will be canned for special treatment, | |
133 | but inspection will not go any deeper than that. |
|
130 | but inspection will not go any deeper than that. | |
134 |
|
131 | |||
135 | Any object whose data is larger than `threshold` will not have their data copied |
|
132 | Any object whose data is larger than `threshold` will not have their data copied | |
136 | (only numpy arrays and bytes/buffers support zero-copy) |
|
133 | (only numpy arrays and bytes/buffers support zero-copy) | |
137 |
|
134 | |||
138 | Message will be a list of bytes/buffers of the format: |
|
135 | Message will be a list of bytes/buffers of the format: | |
139 |
|
136 | |||
140 | [ cf, pinfo, <arg_bufs>, <kwarg_bufs> ] |
|
137 | [ cf, pinfo, <arg_bufs>, <kwarg_bufs> ] | |
141 |
|
138 | |||
142 | With length at least two + len(args) + len(kwargs) |
|
139 | With length at least two + len(args) + len(kwargs) | |
143 | """ |
|
140 | """ | |
144 |
|
141 | |||
145 | arg_bufs = flatten(serialize_object(arg, buffer_threshold, item_threshold) for arg in args) |
|
142 | arg_bufs = flatten(serialize_object(arg, buffer_threshold, item_threshold) for arg in args) | |
146 |
|
143 | |||
147 | kw_keys = sorted(kwargs.keys()) |
|
144 | kw_keys = sorted(kwargs.keys()) | |
148 | kwarg_bufs = flatten(serialize_object(kwargs[key], buffer_threshold, item_threshold) for key in kw_keys) |
|
145 | kwarg_bufs = flatten(serialize_object(kwargs[key], buffer_threshold, item_threshold) for key in kw_keys) | |
149 |
|
146 | |||
150 | info = dict(nargs=len(args), narg_bufs=len(arg_bufs), kw_keys=kw_keys) |
|
147 | info = dict(nargs=len(args), narg_bufs=len(arg_bufs), kw_keys=kw_keys) | |
151 |
|
148 | |||
152 | msg = [pickle.dumps(can(f), PICKLE_PROTOCOL)] |
|
149 | msg = [pickle.dumps(can(f), PICKLE_PROTOCOL)] | |
153 | msg.append(pickle.dumps(info, PICKLE_PROTOCOL)) |
|
150 | msg.append(pickle.dumps(info, PICKLE_PROTOCOL)) | |
154 | msg.extend(arg_bufs) |
|
151 | msg.extend(arg_bufs) | |
155 | msg.extend(kwarg_bufs) |
|
152 | msg.extend(kwarg_bufs) | |
156 |
|
153 | |||
157 | return msg |
|
154 | return msg | |
158 |
|
155 | |||
159 | def unpack_apply_message(bufs, g=None, copy=True): |
|
156 | def unpack_apply_message(bufs, g=None, copy=True): | |
160 | """unpack f,args,kwargs from buffers packed by pack_apply_message() |
|
157 | """unpack f,args,kwargs from buffers packed by pack_apply_message() | |
161 | Returns: original f,args,kwargs""" |
|
158 | Returns: original f,args,kwargs""" | |
162 | bufs = list(bufs) # allow us to pop |
|
159 | bufs = list(bufs) # allow us to pop | |
163 | assert len(bufs) >= 2, "not enough buffers!" |
|
160 | assert len(bufs) >= 2, "not enough buffers!" | |
164 | if not copy: |
|
161 | pf = buffer_to_bytes_py2(bufs.pop(0)) | |
165 | for i in range(2): |
|
162 | f = uncan(pickle.loads(pf), g) | |
166 | bufs[i] = bufs[i].bytes |
|
163 | pinfo = buffer_to_bytes_py2(bufs.pop(0)) | |
167 |
f = |
|
164 | info = pickle.loads(pinfo) | |
168 | info = pickle.loads(bufs.pop(0)) |
|
|||
169 | arg_bufs, kwarg_bufs = bufs[:info['narg_bufs']], bufs[info['narg_bufs']:] |
|
165 | arg_bufs, kwarg_bufs = bufs[:info['narg_bufs']], bufs[info['narg_bufs']:] | |
170 |
|
166 | |||
171 | args = [] |
|
167 | args = [] | |
172 | for i in range(info['nargs']): |
|
168 | for i in range(info['nargs']): | |
173 | arg, arg_bufs = deserialize_object(arg_bufs, g) |
|
169 | arg, arg_bufs = deserialize_object(arg_bufs, g) | |
174 | args.append(arg) |
|
170 | args.append(arg) | |
175 | args = tuple(args) |
|
171 | args = tuple(args) | |
176 | assert not arg_bufs, "Shouldn't be any arg bufs left over" |
|
172 | assert not arg_bufs, "Shouldn't be any arg bufs left over" | |
177 |
|
173 | |||
178 | kwargs = {} |
|
174 | kwargs = {} | |
179 | for key in info['kw_keys']: |
|
175 | for key in info['kw_keys']: | |
180 | kwarg, kwarg_bufs = deserialize_object(kwarg_bufs, g) |
|
176 | kwarg, kwarg_bufs = deserialize_object(kwarg_bufs, g) | |
181 | kwargs[key] = kwarg |
|
177 | kwargs[key] = kwarg | |
182 | assert not kwarg_bufs, "Shouldn't be any kwarg bufs left over" |
|
178 | assert not kwarg_bufs, "Shouldn't be any kwarg bufs left over" | |
183 |
|
179 | |||
184 | return f,args,kwargs |
|
180 | return f,args,kwargs | |
185 |
|
181 |
@@ -1,425 +1,425 b'' | |||||
1 | # encoding: utf-8 |
|
1 | # encoding: utf-8 | |
2 | """Pickle related utilities. Perhaps this should be called 'can'.""" |
|
2 | """Pickle related utilities. Perhaps this should be called 'can'.""" | |
3 |
|
3 | |||
4 | # Copyright (c) IPython Development Team. |
|
4 | # Copyright (c) IPython Development Team. | |
5 | # Distributed under the terms of the Modified BSD License. |
|
5 | # Distributed under the terms of the Modified BSD License. | |
6 |
|
6 | |||
7 | import copy |
|
7 | import copy | |
8 | import logging |
|
8 | import logging | |
9 | import sys |
|
9 | import sys | |
10 | from types import FunctionType |
|
10 | from types import FunctionType | |
11 |
|
11 | |||
12 | try: |
|
12 | try: | |
13 | import cPickle as pickle |
|
13 | import cPickle as pickle | |
14 | except ImportError: |
|
14 | except ImportError: | |
15 | import pickle |
|
15 | import pickle | |
16 |
|
16 | |||
17 | from . import codeutil # This registers a hook when it's imported |
|
17 | from . import codeutil # This registers a hook when it's imported | |
18 | from . import py3compat |
|
18 | from . import py3compat | |
19 | from .importstring import import_item |
|
19 | from .importstring import import_item | |
20 | from .py3compat import string_types, iteritems |
|
20 | from .py3compat import string_types, iteritems, buffer_to_bytes_py2 | |
21 |
|
21 | |||
22 | from IPython.config import Application |
|
22 | from IPython.config import Application | |
23 | from IPython.utils.log import get_logger |
|
23 | from IPython.utils.log import get_logger | |
24 |
|
24 | |||
25 | if py3compat.PY3: |
|
25 | if py3compat.PY3: | |
26 | buffer = memoryview |
|
26 | buffer = memoryview | |
27 | class_type = type |
|
27 | class_type = type | |
28 | else: |
|
28 | else: | |
29 | from types import ClassType |
|
29 | from types import ClassType | |
30 | class_type = (type, ClassType) |
|
30 | class_type = (type, ClassType) | |
31 |
|
31 | |||
32 | try: |
|
32 | try: | |
33 | PICKLE_PROTOCOL = pickle.DEFAULT_PROTOCOL |
|
33 | PICKLE_PROTOCOL = pickle.DEFAULT_PROTOCOL | |
34 | except AttributeError: |
|
34 | except AttributeError: | |
35 | PICKLE_PROTOCOL = pickle.HIGHEST_PROTOCOL |
|
35 | PICKLE_PROTOCOL = pickle.HIGHEST_PROTOCOL | |
36 |
|
36 | |||
37 | def _get_cell_type(a=None): |
|
37 | def _get_cell_type(a=None): | |
38 | """the type of a closure cell doesn't seem to be importable, |
|
38 | """the type of a closure cell doesn't seem to be importable, | |
39 | so just create one |
|
39 | so just create one | |
40 | """ |
|
40 | """ | |
41 | def inner(): |
|
41 | def inner(): | |
42 | return a |
|
42 | return a | |
43 | return type(py3compat.get_closure(inner)[0]) |
|
43 | return type(py3compat.get_closure(inner)[0]) | |
44 |
|
44 | |||
45 | cell_type = _get_cell_type() |
|
45 | cell_type = _get_cell_type() | |
46 |
|
46 | |||
47 | #------------------------------------------------------------------------------- |
|
47 | #------------------------------------------------------------------------------- | |
48 | # Functions |
|
48 | # Functions | |
49 | #------------------------------------------------------------------------------- |
|
49 | #------------------------------------------------------------------------------- | |
50 |
|
50 | |||
51 |
|
51 | |||
52 | def use_dill(): |
|
52 | def use_dill(): | |
53 | """use dill to expand serialization support |
|
53 | """use dill to expand serialization support | |
54 |
|
54 | |||
55 | adds support for object methods and closures to serialization. |
|
55 | adds support for object methods and closures to serialization. | |
56 | """ |
|
56 | """ | |
57 | # import dill causes most of the magic |
|
57 | # import dill causes most of the magic | |
58 | import dill |
|
58 | import dill | |
59 |
|
59 | |||
60 | # dill doesn't work with cPickle, |
|
60 | # dill doesn't work with cPickle, | |
61 | # tell the two relevant modules to use plain pickle |
|
61 | # tell the two relevant modules to use plain pickle | |
62 |
|
62 | |||
63 | global pickle |
|
63 | global pickle | |
64 | pickle = dill |
|
64 | pickle = dill | |
65 |
|
65 | |||
66 | try: |
|
66 | try: | |
67 | from IPython.kernel.zmq import serialize |
|
67 | from IPython.kernel.zmq import serialize | |
68 | except ImportError: |
|
68 | except ImportError: | |
69 | pass |
|
69 | pass | |
70 | else: |
|
70 | else: | |
71 | serialize.pickle = dill |
|
71 | serialize.pickle = dill | |
72 |
|
72 | |||
73 | # disable special function handling, let dill take care of it |
|
73 | # disable special function handling, let dill take care of it | |
74 | can_map.pop(FunctionType, None) |
|
74 | can_map.pop(FunctionType, None) | |
75 |
|
75 | |||
76 | def use_cloudpickle(): |
|
76 | def use_cloudpickle(): | |
77 | """use cloudpickle to expand serialization support |
|
77 | """use cloudpickle to expand serialization support | |
78 |
|
78 | |||
79 | adds support for object methods and closures to serialization. |
|
79 | adds support for object methods and closures to serialization. | |
80 | """ |
|
80 | """ | |
81 | from cloud.serialization import cloudpickle |
|
81 | from cloud.serialization import cloudpickle | |
82 |
|
82 | |||
83 | global pickle |
|
83 | global pickle | |
84 | pickle = cloudpickle |
|
84 | pickle = cloudpickle | |
85 |
|
85 | |||
86 | try: |
|
86 | try: | |
87 | from IPython.kernel.zmq import serialize |
|
87 | from IPython.kernel.zmq import serialize | |
88 | except ImportError: |
|
88 | except ImportError: | |
89 | pass |
|
89 | pass | |
90 | else: |
|
90 | else: | |
91 | serialize.pickle = cloudpickle |
|
91 | serialize.pickle = cloudpickle | |
92 |
|
92 | |||
93 | # disable special function handling, let cloudpickle take care of it |
|
93 | # disable special function handling, let cloudpickle take care of it | |
94 | can_map.pop(FunctionType, None) |
|
94 | can_map.pop(FunctionType, None) | |
95 |
|
95 | |||
96 |
|
96 | |||
97 | #------------------------------------------------------------------------------- |
|
97 | #------------------------------------------------------------------------------- | |
98 | # Classes |
|
98 | # Classes | |
99 | #------------------------------------------------------------------------------- |
|
99 | #------------------------------------------------------------------------------- | |
100 |
|
100 | |||
101 |
|
101 | |||
102 | class CannedObject(object): |
|
102 | class CannedObject(object): | |
103 | def __init__(self, obj, keys=[], hook=None): |
|
103 | def __init__(self, obj, keys=[], hook=None): | |
104 | """can an object for safe pickling |
|
104 | """can an object for safe pickling | |
105 |
|
105 | |||
106 | Parameters |
|
106 | Parameters | |
107 | ========== |
|
107 | ========== | |
108 |
|
108 | |||
109 | obj: |
|
109 | obj: | |
110 | The object to be canned |
|
110 | The object to be canned | |
111 | keys: list (optional) |
|
111 | keys: list (optional) | |
112 | list of attribute names that will be explicitly canned / uncanned |
|
112 | list of attribute names that will be explicitly canned / uncanned | |
113 | hook: callable (optional) |
|
113 | hook: callable (optional) | |
114 | An optional extra callable, |
|
114 | An optional extra callable, | |
115 | which can do additional processing of the uncanned object. |
|
115 | which can do additional processing of the uncanned object. | |
116 |
|
116 | |||
117 | large data may be offloaded into the buffers list, |
|
117 | large data may be offloaded into the buffers list, | |
118 | used for zero-copy transfers. |
|
118 | used for zero-copy transfers. | |
119 | """ |
|
119 | """ | |
120 | self.keys = keys |
|
120 | self.keys = keys | |
121 | self.obj = copy.copy(obj) |
|
121 | self.obj = copy.copy(obj) | |
122 | self.hook = can(hook) |
|
122 | self.hook = can(hook) | |
123 | for key in keys: |
|
123 | for key in keys: | |
124 | setattr(self.obj, key, can(getattr(obj, key))) |
|
124 | setattr(self.obj, key, can(getattr(obj, key))) | |
125 |
|
125 | |||
126 | self.buffers = [] |
|
126 | self.buffers = [] | |
127 |
|
127 | |||
128 | def get_object(self, g=None): |
|
128 | def get_object(self, g=None): | |
129 | if g is None: |
|
129 | if g is None: | |
130 | g = {} |
|
130 | g = {} | |
131 | obj = self.obj |
|
131 | obj = self.obj | |
132 | for key in self.keys: |
|
132 | for key in self.keys: | |
133 | setattr(obj, key, uncan(getattr(obj, key), g)) |
|
133 | setattr(obj, key, uncan(getattr(obj, key), g)) | |
134 |
|
134 | |||
135 | if self.hook: |
|
135 | if self.hook: | |
136 | self.hook = uncan(self.hook, g) |
|
136 | self.hook = uncan(self.hook, g) | |
137 | self.hook(obj, g) |
|
137 | self.hook(obj, g) | |
138 | return self.obj |
|
138 | return self.obj | |
139 |
|
139 | |||
140 |
|
140 | |||
141 | class Reference(CannedObject): |
|
141 | class Reference(CannedObject): | |
142 | """object for wrapping a remote reference by name.""" |
|
142 | """object for wrapping a remote reference by name.""" | |
143 | def __init__(self, name): |
|
143 | def __init__(self, name): | |
144 | if not isinstance(name, string_types): |
|
144 | if not isinstance(name, string_types): | |
145 | raise TypeError("illegal name: %r"%name) |
|
145 | raise TypeError("illegal name: %r"%name) | |
146 | self.name = name |
|
146 | self.name = name | |
147 | self.buffers = [] |
|
147 | self.buffers = [] | |
148 |
|
148 | |||
149 | def __repr__(self): |
|
149 | def __repr__(self): | |
150 | return "<Reference: %r>"%self.name |
|
150 | return "<Reference: %r>"%self.name | |
151 |
|
151 | |||
152 | def get_object(self, g=None): |
|
152 | def get_object(self, g=None): | |
153 | if g is None: |
|
153 | if g is None: | |
154 | g = {} |
|
154 | g = {} | |
155 |
|
155 | |||
156 | return eval(self.name, g) |
|
156 | return eval(self.name, g) | |
157 |
|
157 | |||
158 |
|
158 | |||
159 | class CannedCell(CannedObject): |
|
159 | class CannedCell(CannedObject): | |
160 | """Can a closure cell""" |
|
160 | """Can a closure cell""" | |
161 | def __init__(self, cell): |
|
161 | def __init__(self, cell): | |
162 | self.cell_contents = can(cell.cell_contents) |
|
162 | self.cell_contents = can(cell.cell_contents) | |
163 |
|
163 | |||
164 | def get_object(self, g=None): |
|
164 | def get_object(self, g=None): | |
165 | cell_contents = uncan(self.cell_contents, g) |
|
165 | cell_contents = uncan(self.cell_contents, g) | |
166 | def inner(): |
|
166 | def inner(): | |
167 | return cell_contents |
|
167 | return cell_contents | |
168 | return py3compat.get_closure(inner)[0] |
|
168 | return py3compat.get_closure(inner)[0] | |
169 |
|
169 | |||
170 |
|
170 | |||
171 | class CannedFunction(CannedObject): |
|
171 | class CannedFunction(CannedObject): | |
172 |
|
172 | |||
173 | def __init__(self, f): |
|
173 | def __init__(self, f): | |
174 | self._check_type(f) |
|
174 | self._check_type(f) | |
175 | self.code = f.__code__ |
|
175 | self.code = f.__code__ | |
176 | if f.__defaults__: |
|
176 | if f.__defaults__: | |
177 | self.defaults = [ can(fd) for fd in f.__defaults__ ] |
|
177 | self.defaults = [ can(fd) for fd in f.__defaults__ ] | |
178 | else: |
|
178 | else: | |
179 | self.defaults = None |
|
179 | self.defaults = None | |
180 |
|
180 | |||
181 | closure = py3compat.get_closure(f) |
|
181 | closure = py3compat.get_closure(f) | |
182 | if closure: |
|
182 | if closure: | |
183 | self.closure = tuple( can(cell) for cell in closure ) |
|
183 | self.closure = tuple( can(cell) for cell in closure ) | |
184 | else: |
|
184 | else: | |
185 | self.closure = None |
|
185 | self.closure = None | |
186 |
|
186 | |||
187 | self.module = f.__module__ or '__main__' |
|
187 | self.module = f.__module__ or '__main__' | |
188 | self.__name__ = f.__name__ |
|
188 | self.__name__ = f.__name__ | |
189 | self.buffers = [] |
|
189 | self.buffers = [] | |
190 |
|
190 | |||
191 | def _check_type(self, obj): |
|
191 | def _check_type(self, obj): | |
192 | assert isinstance(obj, FunctionType), "Not a function type" |
|
192 | assert isinstance(obj, FunctionType), "Not a function type" | |
193 |
|
193 | |||
194 | def get_object(self, g=None): |
|
194 | def get_object(self, g=None): | |
195 | # try to load function back into its module: |
|
195 | # try to load function back into its module: | |
196 | if not self.module.startswith('__'): |
|
196 | if not self.module.startswith('__'): | |
197 | __import__(self.module) |
|
197 | __import__(self.module) | |
198 | g = sys.modules[self.module].__dict__ |
|
198 | g = sys.modules[self.module].__dict__ | |
199 |
|
199 | |||
200 | if g is None: |
|
200 | if g is None: | |
201 | g = {} |
|
201 | g = {} | |
202 | if self.defaults: |
|
202 | if self.defaults: | |
203 | defaults = tuple(uncan(cfd, g) for cfd in self.defaults) |
|
203 | defaults = tuple(uncan(cfd, g) for cfd in self.defaults) | |
204 | else: |
|
204 | else: | |
205 | defaults = None |
|
205 | defaults = None | |
206 | if self.closure: |
|
206 | if self.closure: | |
207 | closure = tuple(uncan(cell, g) for cell in self.closure) |
|
207 | closure = tuple(uncan(cell, g) for cell in self.closure) | |
208 | else: |
|
208 | else: | |
209 | closure = None |
|
209 | closure = None | |
210 | newFunc = FunctionType(self.code, g, self.__name__, defaults, closure) |
|
210 | newFunc = FunctionType(self.code, g, self.__name__, defaults, closure) | |
211 | return newFunc |
|
211 | return newFunc | |
212 |
|
212 | |||
213 | class CannedClass(CannedObject): |
|
213 | class CannedClass(CannedObject): | |
214 |
|
214 | |||
215 | def __init__(self, cls): |
|
215 | def __init__(self, cls): | |
216 | self._check_type(cls) |
|
216 | self._check_type(cls) | |
217 | self.name = cls.__name__ |
|
217 | self.name = cls.__name__ | |
218 | self.old_style = not isinstance(cls, type) |
|
218 | self.old_style = not isinstance(cls, type) | |
219 | self._canned_dict = {} |
|
219 | self._canned_dict = {} | |
220 | for k,v in cls.__dict__.items(): |
|
220 | for k,v in cls.__dict__.items(): | |
221 | if k not in ('__weakref__', '__dict__'): |
|
221 | if k not in ('__weakref__', '__dict__'): | |
222 | self._canned_dict[k] = can(v) |
|
222 | self._canned_dict[k] = can(v) | |
223 | if self.old_style: |
|
223 | if self.old_style: | |
224 | mro = [] |
|
224 | mro = [] | |
225 | else: |
|
225 | else: | |
226 | mro = cls.mro() |
|
226 | mro = cls.mro() | |
227 |
|
227 | |||
228 | self.parents = [ can(c) for c in mro[1:] ] |
|
228 | self.parents = [ can(c) for c in mro[1:] ] | |
229 | self.buffers = [] |
|
229 | self.buffers = [] | |
230 |
|
230 | |||
231 | def _check_type(self, obj): |
|
231 | def _check_type(self, obj): | |
232 | assert isinstance(obj, class_type), "Not a class type" |
|
232 | assert isinstance(obj, class_type), "Not a class type" | |
233 |
|
233 | |||
234 | def get_object(self, g=None): |
|
234 | def get_object(self, g=None): | |
235 | parents = tuple(uncan(p, g) for p in self.parents) |
|
235 | parents = tuple(uncan(p, g) for p in self.parents) | |
236 | return type(self.name, parents, uncan_dict(self._canned_dict, g=g)) |
|
236 | return type(self.name, parents, uncan_dict(self._canned_dict, g=g)) | |
237 |
|
237 | |||
238 | class CannedArray(CannedObject): |
|
238 | class CannedArray(CannedObject): | |
239 | def __init__(self, obj): |
|
239 | def __init__(self, obj): | |
240 | from numpy import ascontiguousarray |
|
240 | from numpy import ascontiguousarray | |
241 | self.shape = obj.shape |
|
241 | self.shape = obj.shape | |
242 | self.dtype = obj.dtype.descr if obj.dtype.fields else obj.dtype.str |
|
242 | self.dtype = obj.dtype.descr if obj.dtype.fields else obj.dtype.str | |
243 | self.pickled = False |
|
243 | self.pickled = False | |
244 | if sum(obj.shape) == 0: |
|
244 | if sum(obj.shape) == 0: | |
245 | self.pickled = True |
|
245 | self.pickled = True | |
246 | elif obj.dtype == 'O': |
|
246 | elif obj.dtype == 'O': | |
247 | # can't handle object dtype with buffer approach |
|
247 | # can't handle object dtype with buffer approach | |
248 | self.pickled = True |
|
248 | self.pickled = True | |
249 | elif obj.dtype.fields and any(dt == 'O' for dt,sz in obj.dtype.fields.values()): |
|
249 | elif obj.dtype.fields and any(dt == 'O' for dt,sz in obj.dtype.fields.values()): | |
250 | self.pickled = True |
|
250 | self.pickled = True | |
251 | if self.pickled: |
|
251 | if self.pickled: | |
252 | # just pickle it |
|
252 | # just pickle it | |
253 | self.buffers = [pickle.dumps(obj, PICKLE_PROTOCOL)] |
|
253 | self.buffers = [pickle.dumps(obj, PICKLE_PROTOCOL)] | |
254 | else: |
|
254 | else: | |
255 | # ensure contiguous |
|
255 | # ensure contiguous | |
256 | obj = ascontiguousarray(obj, dtype=None) |
|
256 | obj = ascontiguousarray(obj, dtype=None) | |
257 | self.buffers = [buffer(obj)] |
|
257 | self.buffers = [buffer(obj)] | |
258 |
|
258 | |||
259 | def get_object(self, g=None): |
|
259 | def get_object(self, g=None): | |
260 | from numpy import frombuffer |
|
260 | from numpy import frombuffer | |
261 | data = self.buffers[0] |
|
261 | data = self.buffers[0] | |
262 | if self.pickled: |
|
262 | if self.pickled: | |
263 |
# |
|
263 | # we just pickled it | |
264 | return pickle.loads(data) |
|
264 | return pickle.loads(buffer_to_bytes_py2(data)) | |
265 | else: |
|
265 | else: | |
266 | return frombuffer(data, dtype=self.dtype).reshape(self.shape) |
|
266 | return frombuffer(data, dtype=self.dtype).reshape(self.shape) | |
267 |
|
267 | |||
268 |
|
268 | |||
269 | class CannedBytes(CannedObject): |
|
269 | class CannedBytes(CannedObject): | |
270 | wrap = bytes |
|
270 | wrap = bytes | |
271 | def __init__(self, obj): |
|
271 | def __init__(self, obj): | |
272 | self.buffers = [obj] |
|
272 | self.buffers = [obj] | |
273 |
|
273 | |||
274 | def get_object(self, g=None): |
|
274 | def get_object(self, g=None): | |
275 | data = self.buffers[0] |
|
275 | data = self.buffers[0] | |
276 | return self.wrap(data) |
|
276 | return self.wrap(data) | |
277 |
|
277 | |||
278 | def CannedBuffer(CannedBytes): |
|
278 | def CannedBuffer(CannedBytes): | |
279 | wrap = buffer |
|
279 | wrap = buffer | |
280 |
|
280 | |||
281 | #------------------------------------------------------------------------------- |
|
281 | #------------------------------------------------------------------------------- | |
282 | # Functions |
|
282 | # Functions | |
283 | #------------------------------------------------------------------------------- |
|
283 | #------------------------------------------------------------------------------- | |
284 |
|
284 | |||
285 | def _import_mapping(mapping, original=None): |
|
285 | def _import_mapping(mapping, original=None): | |
286 | """import any string-keys in a type mapping |
|
286 | """import any string-keys in a type mapping | |
287 |
|
287 | |||
288 | """ |
|
288 | """ | |
289 | log = get_logger() |
|
289 | log = get_logger() | |
290 | log.debug("Importing canning map") |
|
290 | log.debug("Importing canning map") | |
291 | for key,value in list(mapping.items()): |
|
291 | for key,value in list(mapping.items()): | |
292 | if isinstance(key, string_types): |
|
292 | if isinstance(key, string_types): | |
293 | try: |
|
293 | try: | |
294 | cls = import_item(key) |
|
294 | cls = import_item(key) | |
295 | except Exception: |
|
295 | except Exception: | |
296 | if original and key not in original: |
|
296 | if original and key not in original: | |
297 | # only message on user-added classes |
|
297 | # only message on user-added classes | |
298 | log.error("canning class not importable: %r", key, exc_info=True) |
|
298 | log.error("canning class not importable: %r", key, exc_info=True) | |
299 | mapping.pop(key) |
|
299 | mapping.pop(key) | |
300 | else: |
|
300 | else: | |
301 | mapping[cls] = mapping.pop(key) |
|
301 | mapping[cls] = mapping.pop(key) | |
302 |
|
302 | |||
303 | def istype(obj, check): |
|
303 | def istype(obj, check): | |
304 | """like isinstance(obj, check), but strict |
|
304 | """like isinstance(obj, check), but strict | |
305 |
|
305 | |||
306 | This won't catch subclasses. |
|
306 | This won't catch subclasses. | |
307 | """ |
|
307 | """ | |
308 | if isinstance(check, tuple): |
|
308 | if isinstance(check, tuple): | |
309 | for cls in check: |
|
309 | for cls in check: | |
310 | if type(obj) is cls: |
|
310 | if type(obj) is cls: | |
311 | return True |
|
311 | return True | |
312 | return False |
|
312 | return False | |
313 | else: |
|
313 | else: | |
314 | return type(obj) is check |
|
314 | return type(obj) is check | |
315 |
|
315 | |||
316 | def can(obj): |
|
316 | def can(obj): | |
317 | """prepare an object for pickling""" |
|
317 | """prepare an object for pickling""" | |
318 |
|
318 | |||
319 | import_needed = False |
|
319 | import_needed = False | |
320 |
|
320 | |||
321 | for cls,canner in iteritems(can_map): |
|
321 | for cls,canner in iteritems(can_map): | |
322 | if isinstance(cls, string_types): |
|
322 | if isinstance(cls, string_types): | |
323 | import_needed = True |
|
323 | import_needed = True | |
324 | break |
|
324 | break | |
325 | elif istype(obj, cls): |
|
325 | elif istype(obj, cls): | |
326 | return canner(obj) |
|
326 | return canner(obj) | |
327 |
|
327 | |||
328 | if import_needed: |
|
328 | if import_needed: | |
329 | # perform can_map imports, then try again |
|
329 | # perform can_map imports, then try again | |
330 | # this will usually only happen once |
|
330 | # this will usually only happen once | |
331 | _import_mapping(can_map, _original_can_map) |
|
331 | _import_mapping(can_map, _original_can_map) | |
332 | return can(obj) |
|
332 | return can(obj) | |
333 |
|
333 | |||
334 | return obj |
|
334 | return obj | |
335 |
|
335 | |||
336 | def can_class(obj): |
|
336 | def can_class(obj): | |
337 | if isinstance(obj, class_type) and obj.__module__ == '__main__': |
|
337 | if isinstance(obj, class_type) and obj.__module__ == '__main__': | |
338 | return CannedClass(obj) |
|
338 | return CannedClass(obj) | |
339 | else: |
|
339 | else: | |
340 | return obj |
|
340 | return obj | |
341 |
|
341 | |||
342 | def can_dict(obj): |
|
342 | def can_dict(obj): | |
343 | """can the *values* of a dict""" |
|
343 | """can the *values* of a dict""" | |
344 | if istype(obj, dict): |
|
344 | if istype(obj, dict): | |
345 | newobj = {} |
|
345 | newobj = {} | |
346 | for k, v in iteritems(obj): |
|
346 | for k, v in iteritems(obj): | |
347 | newobj[k] = can(v) |
|
347 | newobj[k] = can(v) | |
348 | return newobj |
|
348 | return newobj | |
349 | else: |
|
349 | else: | |
350 | return obj |
|
350 | return obj | |
351 |
|
351 | |||
352 | sequence_types = (list, tuple, set) |
|
352 | sequence_types = (list, tuple, set) | |
353 |
|
353 | |||
354 | def can_sequence(obj): |
|
354 | def can_sequence(obj): | |
355 | """can the elements of a sequence""" |
|
355 | """can the elements of a sequence""" | |
356 | if istype(obj, sequence_types): |
|
356 | if istype(obj, sequence_types): | |
357 | t = type(obj) |
|
357 | t = type(obj) | |
358 | return t([can(i) for i in obj]) |
|
358 | return t([can(i) for i in obj]) | |
359 | else: |
|
359 | else: | |
360 | return obj |
|
360 | return obj | |
361 |
|
361 | |||
362 | def uncan(obj, g=None): |
|
362 | def uncan(obj, g=None): | |
363 | """invert canning""" |
|
363 | """invert canning""" | |
364 |
|
364 | |||
365 | import_needed = False |
|
365 | import_needed = False | |
366 | for cls,uncanner in iteritems(uncan_map): |
|
366 | for cls,uncanner in iteritems(uncan_map): | |
367 | if isinstance(cls, string_types): |
|
367 | if isinstance(cls, string_types): | |
368 | import_needed = True |
|
368 | import_needed = True | |
369 | break |
|
369 | break | |
370 | elif isinstance(obj, cls): |
|
370 | elif isinstance(obj, cls): | |
371 | return uncanner(obj, g) |
|
371 | return uncanner(obj, g) | |
372 |
|
372 | |||
373 | if import_needed: |
|
373 | if import_needed: | |
374 | # perform uncan_map imports, then try again |
|
374 | # perform uncan_map imports, then try again | |
375 | # this will usually only happen once |
|
375 | # this will usually only happen once | |
376 | _import_mapping(uncan_map, _original_uncan_map) |
|
376 | _import_mapping(uncan_map, _original_uncan_map) | |
377 | return uncan(obj, g) |
|
377 | return uncan(obj, g) | |
378 |
|
378 | |||
379 | return obj |
|
379 | return obj | |
380 |
|
380 | |||
381 | def uncan_dict(obj, g=None): |
|
381 | def uncan_dict(obj, g=None): | |
382 | if istype(obj, dict): |
|
382 | if istype(obj, dict): | |
383 | newobj = {} |
|
383 | newobj = {} | |
384 | for k, v in iteritems(obj): |
|
384 | for k, v in iteritems(obj): | |
385 | newobj[k] = uncan(v,g) |
|
385 | newobj[k] = uncan(v,g) | |
386 | return newobj |
|
386 | return newobj | |
387 | else: |
|
387 | else: | |
388 | return obj |
|
388 | return obj | |
389 |
|
389 | |||
390 | def uncan_sequence(obj, g=None): |
|
390 | def uncan_sequence(obj, g=None): | |
391 | if istype(obj, sequence_types): |
|
391 | if istype(obj, sequence_types): | |
392 | t = type(obj) |
|
392 | t = type(obj) | |
393 | return t([uncan(i,g) for i in obj]) |
|
393 | return t([uncan(i,g) for i in obj]) | |
394 | else: |
|
394 | else: | |
395 | return obj |
|
395 | return obj | |
396 |
|
396 | |||
397 | def _uncan_dependent_hook(dep, g=None): |
|
397 | def _uncan_dependent_hook(dep, g=None): | |
398 | dep.check_dependency() |
|
398 | dep.check_dependency() | |
399 |
|
399 | |||
400 | def can_dependent(obj): |
|
400 | def can_dependent(obj): | |
401 | return CannedObject(obj, keys=('f', 'df'), hook=_uncan_dependent_hook) |
|
401 | return CannedObject(obj, keys=('f', 'df'), hook=_uncan_dependent_hook) | |
402 |
|
402 | |||
403 | #------------------------------------------------------------------------------- |
|
403 | #------------------------------------------------------------------------------- | |
404 | # API dictionaries |
|
404 | # API dictionaries | |
405 | #------------------------------------------------------------------------------- |
|
405 | #------------------------------------------------------------------------------- | |
406 |
|
406 | |||
407 | # These dicts can be extended for custom serialization of new objects |
|
407 | # These dicts can be extended for custom serialization of new objects | |
408 |
|
408 | |||
409 | can_map = { |
|
409 | can_map = { | |
410 | 'IPython.parallel.dependent' : can_dependent, |
|
410 | 'IPython.parallel.dependent' : can_dependent, | |
411 | 'numpy.ndarray' : CannedArray, |
|
411 | 'numpy.ndarray' : CannedArray, | |
412 | FunctionType : CannedFunction, |
|
412 | FunctionType : CannedFunction, | |
413 | bytes : CannedBytes, |
|
413 | bytes : CannedBytes, | |
414 | buffer : CannedBuffer, |
|
414 | buffer : CannedBuffer, | |
415 | cell_type : CannedCell, |
|
415 | cell_type : CannedCell, | |
416 | class_type : can_class, |
|
416 | class_type : can_class, | |
417 | } |
|
417 | } | |
418 |
|
418 | |||
419 | uncan_map = { |
|
419 | uncan_map = { | |
420 | CannedObject : lambda obj, g: obj.get_object(g), |
|
420 | CannedObject : lambda obj, g: obj.get_object(g), | |
421 | } |
|
421 | } | |
422 |
|
422 | |||
423 | # for use in _import_mapping: |
|
423 | # for use in _import_mapping: | |
424 | _original_can_map = can_map.copy() |
|
424 | _original_can_map = can_map.copy() | |
425 | _original_uncan_map = uncan_map.copy() |
|
425 | _original_uncan_map = uncan_map.copy() |
General Comments 0
You need to be logged in to leave comments.
Login now