code.kerkeslager.com Git - sandbox/blob - serial/binary.py

   1 import collections
   2 import functools
   3 import io
   4 import struct
   5
   6 TAG_NULL = 0x00
   7 TAG_TRUE = 0x01
   8 TAG_FALSE = 0x02
   9 TAG_UINT8 = 0x03
  10 TAG_UINT16 = 0x04
  11 TAG_UINT32 = 0x05
  12 TAG_UINT64 = 0x06
  13 TAG_INT8 = 0x10
  14 TAG_INT16 = 0x11
  15 TAG_INT32 = 0x12
  16 TAG_INT64 = 0x13
  17 TAG_BINARY = 0x20
  18 TAG_UTF8 = 0x21
  19 TAG_UTF16 = 0x22
  20 TAG_UTF32 = 0x23
  21 TAG_TUPLE = 0x30
  22
  23 TaggedObject = collections.namedtuple(
  24     'TaggedObject',
  25     [
  26         'tag',
  27         'instance',
  28     ],
  29 )
  30
  31 def _make_tag_only_serializer(tag, expected_value):
  32     tag = bytes([tag])
  33
  34     def serializer(to):
  35         assert to.instance == expected_value
  36         return tag
  37
  38     return serializer
  39
  40 def _make_struct_serializer(fmt):
  41     fmt = '!B' + fmt
  42     packer = functools.partial(struct.pack, fmt)
  43
  44     def serializer(to):
  45         return packer(to.tag, to.instance)
  46
  47     return serializer
  48
  49 def _make_string_serializer(encoder):
  50     packer = functools.partial(struct.pack, '!BI')
  51
  52     def serializer(to):
  53         encoded = encoder(to.instance)
  54         return packer(to.tag, len(encoded)) + encoded
  55
  56     return serializer
  57
  58 def _serialize_tuple(to):
  59     assert isinstance(to.instance, tuple)
  60
  61     payload = b''.join(serialize(item) for item in to.instance)
  62
  63     fmt = '!BI'
  64
  65     return struct.pack('!BI', TAG_TUPLE, len(payload)) + payload
  66
  67
  68 _TAGS_TO_SERIALIZERS = {
  69     TAG_NULL: _make_tag_only_serializer(TAG_NULL, None),
  70     TAG_TRUE: _make_tag_only_serializer(TAG_TRUE, True),
  71     TAG_FALSE: _make_tag_only_serializer(TAG_FALSE, False),
  72     TAG_UINT8: _make_struct_serializer('B'),
  73     TAG_UINT16: _make_struct_serializer('H'),
  74     TAG_UINT32: _make_struct_serializer('I'),
  75     TAG_UINT64: _make_struct_serializer('Q'),
  76     TAG_INT8: _make_struct_serializer('b'),
  77     TAG_INT16: _make_struct_serializer('h'),
  78     TAG_INT32: _make_struct_serializer('i'),
  79     TAG_INT64: _make_struct_serializer('q'),
  80     TAG_BINARY: _make_string_serializer(lambda s: s),
  81     TAG_UTF8: _make_string_serializer(lambda s: s.encode('utf-8')),
  82     TAG_UTF16: _make_string_serializer(lambda s: s.encode('utf-16')),
  83     TAG_UTF32: _make_string_serializer(lambda s: s.encode('utf-32')),
  84     TAG_TUPLE: _serialize_tuple,
  85 }
  86
  87 def serialize(to):
  88     return _TAGS_TO_SERIALIZERS[to.tag](to)
  89
  90 def _make_tag_only_parser(tag, value):
  91     def parser(b):
  92         return 0, TaggedObject(tag = tag, instance = value)
  93
  94     return parser
  95
  96 def _make_struct_deserializer(tag, fmt):
  97     fmt = '!' + fmt
  98     size = struct.calcsize(fmt)
  99     unpacker = functools.partial(struct.unpack, fmt)
 100
 101     def parser(b):
 102         b = b.read(size)
 103         assert len(b) == size
 104         return size, TaggedObject(tag = tag, instance = unpacker(b)[0])
 105
 106     return parser
 107
 108 _LENGTH_FMT = '!I'
 109 _LENGTH_FMT_SIZE = struct.calcsize(_LENGTH_FMT)
 110
 111 def _read_length_then_payload(b):
 112     length_b = b.read(_LENGTH_FMT_SIZE)
 113     assert len(length_b) == _LENGTH_FMT_SIZE
 114     length = struct.unpack(_LENGTH_FMT, length_b)[0]
 115
 116     payload = b.read(length)
 117     assert len(payload) == length
 118     return _LENGTH_FMT_SIZE + length, payload
 119
 120 def _make_string_deserializer(tag, decoder):
 121     fmt = '!I'
 122     size = struct.calcsize(fmt)
 123     unpacker = functools.partial(struct.unpack, fmt)
 124
 125     def parser(b):
 126         bytes_read, payload = _read_length_then_payload(b)
 127         return bytes_read, TaggedObject(tag = tag, instance = decoder(payload))
 128
 129     return parser
 130
 131 def _deserialize_tuple(b):
 132     bytes_read, payload = _read_length_then_payload(b)
 133
 134     payload_stream = io.BytesIO(payload)
 135
 136     total_bytes_read = 0
 137     instance = []
 138
 139     while total_bytes_read < len(payload):
 140         partial_bytes_read, item = _deserialize_partial(payload_stream)
 141         total_bytes_read += partial_bytes_read
 142         instance.append(item)
 143
 144     return bytes_read, TaggedObject(tag = TAG_TUPLE, instance = tuple(instance))
 145
 146 _TAGS_TO_PARSERS = {
 147     TAG_NULL: _make_tag_only_parser(TAG_NULL, None),
 148     TAG_TRUE: _make_tag_only_parser(TAG_TRUE, True),
 149     TAG_FALSE: _make_tag_only_parser(TAG_FALSE, False),
 150     TAG_UINT8: _make_struct_deserializer(TAG_UINT8, 'B'),
 151     TAG_UINT16: _make_struct_deserializer(TAG_UINT16, 'H'),
 152     TAG_UINT32: _make_struct_deserializer(TAG_UINT32, 'I'),
 153     TAG_UINT64: _make_struct_deserializer(TAG_UINT64, 'Q'),
 154     TAG_INT8: _make_struct_deserializer(TAG_INT8, 'b'),
 155     TAG_INT16: _make_struct_deserializer(TAG_INT16, 'h'),
 156     TAG_INT32: _make_struct_deserializer(TAG_INT32, 'i'),
 157     TAG_INT64: _make_struct_deserializer(TAG_INT64, 'q'),
 158     TAG_BINARY: _make_string_deserializer(TAG_BINARY, lambda b: b),
 159     TAG_UTF8: _make_string_deserializer(TAG_UTF8, lambda b: b.decode('utf-8')),
 160     TAG_UTF16: _make_string_deserializer(TAG_UTF16, lambda b: b.decode('utf-16')),
 161     TAG_UTF32: _make_string_deserializer(TAG_UTF32, lambda b: b.decode('utf-32')),
 162     TAG_TUPLE: _deserialize_tuple,
 163 }
 164
 165 def _deserialize_partial(b):
 166     tag = b.read(1)
 167     assert len(tag) == 1
 168     bytes_read, to = _TAGS_TO_PARSERS[tag[0]](b)
 169     return bytes_read + 1, to
 170
 171 def deserialize(b):
 172     if isinstance(b, bytes):
 173         b = io.BytesIO(b)
 174
 175     bytes_read, result = _deserialize_partial(b)
 176
 177     remainder = b.read()
 178
 179     if len(remainder) == 0:
 180         return result
 181
 182     raise Exception('Unable to parse remainder: {}'.format(remainder))