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

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