pickle_zero_copy.py

"""
Simple proof-of-concept for zero-copy unpickling
"""
import pickle
import struct


def serialize(obj):
    buffers = []
    pickled = pickle.dumps(obj, protocol=pickle.HIGHEST_PROTOCOL, buffer_callback=buffers.append)
    buffers = [buffer.raw() for buffer in buffers]
    buffers.append(pickled)
    return serialize_buffers(buffers)


def deserialize(data):
    buffers = list(deserialize_buffers(data))
    return pickle.loads(buffers[-1], buffers=buffers)


def serialize_buffers(buffers):
    sizes = [len(buffer) for buffer in buffers]
    header = serialize_ints(sizes)
    return header + b''.join(buffers)


def deserialize_buffers(data):
    data = memoryview(data)
    sizes, offset = deserialize_ints(data)
    for size in sizes:
        next_offset = offset + size
        yield data[offset:next_offset]
        offset = next_offset


def serialize_ints(ints):
    count = len(ints)
    return struct.pack(f'<{count + 1}I', count, *ints)


def deserialize_ints(data):
    count = struct.unpack('<I', data[:4])[0]
    end = 4 + (count * 4)
    return struct.unpack(f'<{count}I', data[4:end]), end