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

   1 import functools
   2 import io
   3 import struct
   4
   5 from . import tags
   6
   7 def _make_tag_only_serializer(tag, expected_value):
   8     tag = bytes([tag])
   9
  10     def serializer(to):
  11         assert to.instance == expected_value
  12         return tag
  13
  14     return serializer
  15
  16 def _make_struct_serializer(fmt):
  17     fmt = '!B' + fmt
  18     packer = functools.partial(struct.pack, fmt)
  19
  20     def serializer(to):
  21         return packer(to.tag, to.instance)
  22
  23     return serializer
  24
  25 def _make_string_serializer(encoder):
  26     packer = functools.partial(struct.pack, '!BI')
  27
  28     def serializer(to):
  29         encoded = encoder(to.instance)
  30         return packer(to.tag, len(encoded)) + encoded
  31
  32     return serializer
  33
  34 def _serialize_tuple(to):
  35     assert isinstance(to.instance, tuple)
  36
  37     payload = b''.join(serialize(item) for item in to.instance)
  38
  39     fmt = '!BI'
  40
  41     return struct.pack('!BI', tags.TUPLE, len(payload)) + payload
  42
  43 def _serialize_list(to):
  44     assert isinstance(to.instance, list)
  45
  46     # TODO Actually handle this case somehow
  47     assert len(to.instance) > 0
  48
  49     # TODO Do this a better way
  50     serialized_items = [serialize(i) for i in to.instance]
  51     list_tag = serialized_items[0][0]
  52
  53     def check_and_strip_prefix(b):
  54         item_tag = b[0]
  55         assert list_tag == item_tag
  56         return b[1:]
  57
  58     payload = b''.join(check_and_strip_prefix(si) for si in serialized_items)
  59
  60     fmt = '!BBI'
  61
  62     return struct.pack(fmt, tags.LIST, list_tag, len(payload)) + payload
  63
  64 def _serialize_object(to):
  65     raise Exception('Not implemented')
  66
  67 _TAGS_TO_SERIALIZERS = {
  68     tags.NULL: _make_tag_only_serializer(tags.NULL, None),
  69     tags.TRUE: _make_tag_only_serializer(tags.TRUE, True),
  70     tags.FALSE: _make_tag_only_serializer(tags.FALSE, False),
  71     tags.UINT8: _make_struct_serializer('B'),
  72     tags.UINT16: _make_struct_serializer('H'),
  73     tags.UINT32: _make_struct_serializer('I'),
  74     tags.UINT64: _make_struct_serializer('Q'),
  75     tags.INT8: _make_struct_serializer('b'),
  76     tags.INT16: _make_struct_serializer('h'),
  77     tags.INT32: _make_struct_serializer('i'),
  78     tags.INT64: _make_struct_serializer('q'),
  79     tags.BINARY: _make_string_serializer(lambda s: s),
  80     tags.UTF8: _make_string_serializer(lambda s: s.encode('utf-8')),
  81     tags.UTF16: _make_string_serializer(lambda s: s.encode('utf-16')),
  82     tags.UTF32: _make_string_serializer(lambda s: s.encode('utf-32')),
  83     tags.TUPLE: _serialize_tuple,
  84     tags.LIST: _serialize_list,
  85     tags.OBJECT: _serialize_object,
  86 }
  87
  88 def serialize(to):
  89     return _TAGS_TO_SERIALIZERS[to.tag](to)
  90
  91 def _make_tag_only_parser(tag, value):
  92     def parser(b):
  93         return 0, tags.TaggedObject(tag = tag, instance = value)
  94
  95     return parser
  96
  97 def _make_struct_deserializer(tag, fmt):
  98     fmt = '!' + fmt
  99     size = struct.calcsize(fmt)
 100     unpacker = functools.partial(struct.unpack, fmt)
 101
 102     def parser(b):
 103         b = b.read(size)
 104         assert len(b) == size
 105         return size, tags.TaggedObject(tag = tag, instance = unpacker(b)[0])
 106
 107     return parser
 108
 109 _LENGTH_FMT = '!I'
 110 _LENGTH_FMT_SIZE = struct.calcsize(_LENGTH_FMT)
 111
 112 def _read_length_then_payload(b):
 113     length_b = b.read(_LENGTH_FMT_SIZE)
 114     assert len(length_b) == _LENGTH_FMT_SIZE
 115     length = struct.unpack(_LENGTH_FMT, length_b)[0]
 116
 117     payload = b.read(length)
 118     assert len(payload) == length
 119     return _LENGTH_FMT_SIZE + length, payload
 120
 121 def _make_string_deserializer(tag, decoder):
 122     fmt = '!I'
 123     size = struct.calcsize(fmt)
 124     unpacker = functools.partial(struct.unpack, fmt)
 125
 126     def parser(b):
 127         bytes_read, payload = _read_length_then_payload(b)
 128         return bytes_read, tags.TaggedObject(tag = tag, instance = decoder(payload))
 129
 130     return parser
 131
 132 def _deserialize_tuple(b):
 133     bytes_read, payload = _read_length_then_payload(b)
 134
 135     payload_stream = io.BytesIO(payload)
 136
 137     total_bytes_read = 0
 138     instance = []
 139
 140     while total_bytes_read < len(payload):
 141         partial_bytes_read, item = _deserialize_partial(payload_stream)
 142         total_bytes_read += partial_bytes_read
 143         instance.append(item)
 144
 145     return bytes_read, tags.TaggedObject(tag = tags.TUPLE, instance = tuple(instance))
 146
 147 def _deserialize_list(b):
 148     list_tag_bytes = b.read(1)
 149     assert len(list_tag_bytes) == 1
 150     list_tag = list_tag_bytes[0]
 151
 152     bytes_read, payload = _read_length_then_payload(b)
 153
 154     payload_stream = io.BytesIO(payload)
 155
 156     total_bytes_read = 0
 157     instance = []
 158
 159     while total_bytes_read < len(payload):
 160         partial_bytes_read, item = _TAGS_TO_PARSERS[list_tag](payload_stream)
 161         total_bytes_read += partial_bytes_read
 162         instance.append(item)
 163
 164     # TODO Return tags = (tags.LIST, list_tag) to function like a generic type
 165     return bytes_read, tags.TaggedObject(tag = tags.LIST, instance = instance)
 166
 167 def _deserialize_object(b):
 168     raise Exception('Not implemented')
 169
 170 _TAGS_TO_PARSERS = {
 171     tags.NULL: _make_tag_only_parser(tags.NULL, None),
 172     tags.TRUE: _make_tag_only_parser(tags.TRUE, True),
 173     tags.FALSE: _make_tag_only_parser(tags.FALSE, False),
 174     tags.UINT8: _make_struct_deserializer(tags.UINT8, 'B'),
 175     tags.UINT16: _make_struct_deserializer(tags.UINT16, 'H'),
 176     tags.UINT32: _make_struct_deserializer(tags.UINT32, 'I'),
 177     tags.UINT64: _make_struct_deserializer(tags.UINT64, 'Q'),
 178     tags.INT8: _make_struct_deserializer(tags.INT8, 'b'),
 179     tags.INT16: _make_struct_deserializer(tags.INT16, 'h'),
 180     tags.INT32: _make_struct_deserializer(tags.INT32, 'i'),
 181     tags.INT64: _make_struct_deserializer(tags.INT64, 'q'),
 182     tags.BINARY: _make_string_deserializer(tags.BINARY, lambda b: b),
 183     tags.UTF8: _make_string_deserializer(tags.UTF8, lambda b: b.decode('utf-8')),
 184     tags.UTF16: _make_string_deserializer(tags.UTF16, lambda b: b.decode('utf-16')),
 185     tags.UTF32: _make_string_deserializer(tags.UTF32, lambda b: b.decode('utf-32')),
 186     tags.TUPLE: _deserialize_tuple,
 187     tags.LIST: _deserialize_list,
 188     tags.OBJECT: _deserialize_object,
 189 }
 190
 191 def _deserialize_partial(b):
 192     tag = b.read(1)
 193     assert len(tag) == 1
 194     bytes_read, to = _TAGS_TO_PARSERS[tag[0]](b)
 195     return bytes_read + 1, to
 196
 197 def deserialize(b):
 198     if isinstance(b, bytes):
 199         b = io.BytesIO(b)
 200
 201     bytes_read, result = _deserialize_partial(b)
 202
 203     remainder = b.read()
 204
 205     if len(remainder) == 0:
 206         return result
 207
 208     raise Exception('Unable to parse remainder: {}'.format(remainder))