A simple version of pickle

simplepickle.py

import pickle
import struct


class Memo(object):
    def __init__(self):
        self.kv = {}

    def put(self, item):
        ind = self.kv[id(item)] = len(self.kv)
        if ind >= 256:
            raise NotImplementedError("memos longer than 256")
        return ind

    def get(self, item):
        return self.kv.get(id(item))


def dumps(x):
    """A simple version of pickle.dumps, only supporting a few types"""
    buf = bytearray()
    memo = Memo()

    # All pickles start with PROTO + VERSION
    buf.extend(pickle.PROTO)
    buf.extend(b"\x03")

    # Recurively pickle the contents of x
    _simple_pickle(x, buf, memo)

    # All pickles end with STOP
    buf.extend(pickle.STOP)

    return bytes(buf)


def _simple_pickle(x, buf, memo):
    if x is None:
        # Things like None/True/False serialize as single bytes in the protocol
        buf.extend(pickle.NONE)
    elif x is True:
        buf.extend(pickle.NEWTRUE)
    elif x is False:
        buf.extend(pickle.NEWFALSE)
    elif isinstance(x, int):
        # Integers have a few different storage forms dependent on size
        # we'll only support one of these here.
        if 0 <= x < 256:
            # Small integers serialize as an opcode followed by
            # the integers stored as an unsigned 1byte int
            buf.extend(pickle.BININT1)
            buf.extend(struct.pack("!B", x))
        else:
            raise NotImplementedError("ints not in 0 <= ind < 255")
    elif isinstance(x, list):
        if memo.get(x) is not None:
            # We've already seen this object, get it from the memo
            # instead of reserializing.
            buf.extend(pickle.BINGET)
            buf.extend(struct.pack("!B", memo.get(x)))
        else:
            # Lists are built by:
            # - pushing an empty list to the stack,
            # - pushing their contents to the stack,
            # - then bulk appending everything in between the last
            #   MARK and an APPENDS flag
            buf.extend(pickle.EMPTY_LIST)
            # Lists aren't singletons (like everything above), so
            # we need to maintain references to them. Store the list
            # in the memo to possibly be used later.
            ind = memo.put(x)
            buf.extend(pickle.BINPUT)
            buf.extend(struct.pack("!B", ind))
            # Pickle the elements of the list
            buf.extend(pickle.MARK)
            for i in x:
                _simple_pickle(i, buf, memo)
            # Bulk append all the elements
            buf.extend(pickle.APPENDS)
    else:
        raise NotImplementedError