def pack_apply_message(f, args, kwargs, buffer_threshold=MAX_BYTES, item_threshold=MAX_ITEMS): """pack up a function, args, and kwargs to be sent over the wire Each element of args/kwargs will be canned for special treatment, but inspection will not go any deeper than that. Any object whose data is larger than `threshold` will not have their data copied (only numpy arrays and bytes/buffers support zero-copy) Message will be a list of bytes/buffers of the format: [ cf, pinfo, <arg_bufs>, <kwarg_bufs> ] With length at least two + len(args) + len(kwargs) """ arg_bufs = list( chain.from_iterable(serialize_object(arg, buffer_threshold, item_threshold) for arg in args) ) kw_keys = sorted(kwargs.keys()) kwarg_bufs = list( chain.from_iterable( serialize_object(kwargs[key], buffer_threshold, item_threshold) for key in kw_keys ) ) info = dict(nargs=len(args), narg_bufs=len(arg_bufs), kw_keys=kw_keys) msg = [pickle.dumps(can(f), PICKLE_PROTOCOL)] msg.append(pickle.dumps(info, PICKLE_PROTOCOL)) msg.extend(arg_bufs) msg.extend(kwarg_bufs) return msg
def test_can_partial(): def foo(x, y, z): return x * y * z partial_foo = partial(foo, 2, y=5) canned = can(partial_foo) assert isinstance(canned, canning.CannedPartial) dumped = pickle.dumps(canned) loaded = pickle.loads(dumped) pfoo2 = uncan(loaded) assert pfoo2(z=3) == partial_foo(z=3)
def serialize_object(obj, buffer_threshold=MAX_BYTES, item_threshold=MAX_ITEMS): """Serialize an object into a list of sendable buffers. Parameters ---------- obj : object The object to be serialized buffer_threshold : int The threshold (in bytes) for pulling out data buffers to avoid pickling them. item_threshold : int The maximum number of items over which canning will iterate. Containers (lists, dicts) larger than this will be pickled without introspection. Returns ------- [bufs] : list of buffers representing the serialized object. """ buffers = [] if istype(obj, sequence_types) and len(obj) < item_threshold: cobj = can_sequence(obj) for c in cobj: buffers.extend(_extract_buffers(c, buffer_threshold)) elif istype(obj, dict) and len(obj) < item_threshold: cobj = {} for k in sorted(obj): c = can(obj[k]) buffers.extend(_extract_buffers(c, buffer_threshold)) cobj[k] = c else: cobj = can(obj) buffers.extend(_extract_buffers(cobj, buffer_threshold)) buffers.insert(0, pickle.dumps(cobj, PICKLE_PROTOCOL)) return buffers
def test_can_partial_buffers(): def foo(arg1, arg2, kwarg1, kwarg2): return '%s%s%s%s' % (arg1, arg2[:32], b2a_hex(kwarg1.tobytes()[:32]), kwarg2) buf1 = os.urandom(1024 * 1024) buf2 = memoryview(os.urandom(1024 * 1024)) partial_foo = partial(foo, 5, buf1, kwarg1=buf2, kwarg2=10) canned = can(partial_foo) assert len(canned.buffers) == 2 assert isinstance(canned, canning.CannedPartial) buffers, canned.buffers = canned.buffers, [] dumped = pickle.dumps(canned) loaded = pickle.loads(dumped) loaded.buffers = buffers pfoo2 = uncan(loaded) assert pfoo2() == partial_foo()
def test_can_partial_buffers(): def foo(arg1, arg2, kwarg1, kwarg2): return f'{arg1}{arg2[:32]}{b2a_hex(kwarg1.tobytes()[:32])}{kwarg2}' buf1 = os.urandom(1024 * 1024) buf2 = memoryview(os.urandom(1024 * 1024)) partial_foo = partial(foo, 5, buf1, kwarg1=buf2, kwarg2=10) canned = can(partial_foo) assert len(canned.buffers) == 2 assert isinstance(canned, canning.CannedPartial) buffers, canned.buffers = canned.buffers, [] dumped = pickle.dumps(canned) loaded = pickle.loads(dumped) loaded.buffers = buffers pfoo2 = uncan(loaded) assert pfoo2() == partial_foo()
def test_uncan_bytes_buffer(): data = b'data' canned = can(data) canned.buffers = [memoryview(buf) for buf in canned.buffers] out = uncan(canned) assert out == data
def dumps(obj): return pickle.dumps(can(obj))
def test_uncan_bytes_buffer(): data = b'data' canned = can(data) canned.buffers = [memoryview(buf) for buf in canned.buffers] out = uncan(canned) nt.assert_equal(out, data)