code.kerkeslager.com Git - ton/blob - don/__init__.py

   1 import collections
   2 import struct
   3
   4 VOID = 0x00
   5 TRUE = 0x01
   6 FALSE = 0x02
   7 BOOL = (TRUE, FALSE)
   8 INT8 = 0x10
   9 INT16 = 0x11
  10 INT32 = 0x12
  11 INT64 = 0x13
  12 FLOAT = 0x20
  13 DOUBLE = 0x21
  14 BINARY = 0x30
  15 UTF8 = 0x31
  16 UTF16 = 0x32
  17 UTF32 = 0x33
  18 LIST = 0x40
  19 DICTIONARY = 0x41
  20
  21 DEFAULT_INTEGER_ENCODING = INT32
  22 DEFAULT_DECIMAL_ENCODING = DOUBLE
  23 DEFAULT_STRING_ENCODING = UTF8
  24
  25 TaggedObject = collections.namedtuple('TaggedObject', ['tag', 'value'])
  26
  27 _TYPES_TO_TAGS = {
  28     int: DEFAULT_INTEGER_ENCODING,
  29     float: DEFAULT_DECIMAL_ENCODING,
  30     bytes: BINARY,
  31     str: DEFAULT_STRING_ENCODING,
  32     list: LIST,
  33     dict: DICTIONARY,
  34     collections.OrderedDict: DICTIONARY,
  35 }
  36
  37 def tag(o):
  38     if isinstance(o, TaggedObject):
  39         return o
  40
  41     if o is None:
  42         return TaggedObject(tag = VOID, value = o)
  43
  44     if o is True:
  45         return TaggedObject(tag = TRUE, value = o)
  46
  47     if o is False:
  48         return TaggedObject(tag = FALSE, value = o)
  49
  50     return TaggedObject(tag = _TYPES_TO_TAGS[type(o)], value = o)
  51
  52 def serialize_tag_only_type(o):
  53     return b''
  54
  55 def make_serializer_from_pack_format_string(pfs):
  56     def serializer(i):
  57         return struct.pack(pfs, i)
  58     return serializer
  59
  60 def make_string_serializer_from_encoder(e):
  61     def serializer(s):
  62         encoded = e(s)
  63         return struct.pack('!I', len(encoded)) + encoded
  64     return serializer
  65
  66 def serialize_list(items):
  67     # TODO Enforce that items are all the same type
  68     items = [tag(i) for i in items]
  69
  70     if len(items) == 0:
  71         item_tag = VOID
  72     else:
  73         item_tag = items[0].tag
  74
  75     item_serializer = _SERIALIZERS[item_tag]
  76     items = [item_serializer(i.value) for i in items]
  77     item_length = len(items)
  78     items = b''.join(items)
  79     byte_length = len(items)
  80     return struct.pack('!BII', item_tag, byte_length, item_length) + items
  81
  82 def serialize_dict(d):
  83     item_length = 0
  84     serialized = b''
  85
  86     key_serializer = _SERIALIZERS[UTF8]
  87
  88     for key, value in d.items():
  89         assert isinstance(key, str)
  90         item_length += 1
  91         serialized += key_serializer(key) + _binary_serialize(value)
  92
  93     byte_length = len(serialized)
  94     return struct.pack('!II', byte_length, item_length) + serialized
  95
  96 _SERIALIZERS = {
  97     VOID: serialize_tag_only_type,
  98     TRUE: serialize_tag_only_type,
  99     FALSE: serialize_tag_only_type,
 100     INT8: make_serializer_from_pack_format_string('!b'),
 101     INT16: make_serializer_from_pack_format_string('!h'),
 102     INT32: make_serializer_from_pack_format_string('!i'),
 103     FLOAT: make_serializer_from_pack_format_string('!f'),
 104     DOUBLE: make_serializer_from_pack_format_string('!d'),
 105     BINARY: make_string_serializer_from_encoder(lambda b: b),
 106     UTF8: make_string_serializer_from_encoder(lambda s: s.encode('utf-8')),
 107     UTF16: make_string_serializer_from_encoder(lambda s: s.encode('utf-16')),
 108     UTF32: make_string_serializer_from_encoder(lambda s: s.encode('utf-32')),
 109     LIST: serialize_list,
 110     DICTIONARY: serialize_dict,
 111 }
 112
 113 def _binary_serialize(o):
 114     o = tag(o)
 115     return struct.pack('!B', o.tag) + _SERIALIZERS[o.tag](o.value)
 116
 117 ParseResult = collections.namedtuple(
 118     'ParseResult',
 119     [
 120         'success',
 121         'value',
 122         'remaining',
 123     ],
 124 )
 125
 126 _FAILED_PARSE_RESULT = ParseResult(success = False, value = None, remaining = None)
 127
 128 _BYTE_SIZES_TO_UNPACK_FORMATS = {
 129     1: '!b',
 130     2: '!h',
 131     4: '!i',
 132     8: '!q',
 133 }
 134
 135 def make_integer_parser(size_in_bytes):
 136     unpack_format = _BYTE_SIZES_TO_UNPACK_FORMATS[size_in_bytes]
 137
 138     def integer_parser(source):
 139         value = struct.unpack(unpack_format, source[:size_in_bytes])[0]
 140         remaining = source[size_in_bytes:]
 141
 142         return ParseResult(success = True, value = value, remaining = remaining)
 143
 144     return integer_parser
 145
 146 def binary64_parser(source):
 147     return ParseResult(
 148         success = True,
 149         value = struct.unpack('!d', source[:8])[0],
 150         remaining = source[8:],
 151     )
 152
 153 def make_string_parser(decoder):
 154     def string_parser(source):
 155         length = struct.unpack('!I', source[:4])[0]
 156         source = source[4:]
 157         return ParseResult(
 158             success = True,
 159             value = decoder(source[:length]),
 160             remaining = source[length:],
 161         )
 162
 163     return string_parser
 164
 165 def list_parser(source):
 166     tag = source[0]
 167     parser = _TAGS_TO_PARSERS[tag]
 168
 169     source = source[1:]
 170     byte_length, items_length = struct.unpack('!II', source[:8])
 171     source = source[8:]
 172
 173     remaining = source[byte_length:]
 174     source = source[:byte_length]
 175
 176     def item_iterator(source):
 177         count = 0
 178
 179         while len(source) > 0:
 180             parse_result = parser(source)
 181
 182             if parse_result.success:
 183                 count += 1
 184                 yield parse_result.value
 185                 source = parse_result.remaining
 186
 187         assert count == items_length
 188
 189     return ParseResult(
 190         success = True,
 191         value = item_iterator(source),
 192         remaining = remaining,
 193     )
 194
 195 def dictionary_parser(source):
 196     key_parser = _TAGS_TO_PARSERS[UTF8]
 197
 198     byte_length, item_length = struct.unpack('!II', source[:8])
 199     source = source[8:]
 200
 201     remaining = source[byte_length:]
 202     source = source[:byte_length]
 203
 204     def kvp_iterator(source):
 205         count = 0
 206
 207         while len(source) > 0:
 208             count += 1
 209             key_parse_result = key_parser(source)
 210             key, source = key_parse_result.value, key_parse_result.remaining
 211             value_parse_result = _object_parser(source)
 212             value, source = value_parse_result.value, value_parse_result.remaining
 213
 214             yield key, value
 215
 216         assert count == item_length
 217
 218     return ParseResult(
 219         success = True,
 220         value = collections.OrderedDict(kvp_iterator(source)),
 221         remaining = remaining,
 222     )
 223
 224
 225 _TAGS_TO_PARSERS = {
 226     VOID: lambda r: ParseResult(True, None, r),
 227     TRUE: lambda r: ParseResult(True, True, r),
 228     FALSE: lambda r: ParseResult(True, False, r),
 229     INT8: make_integer_parser(1),
 230     INT16: make_integer_parser(2),
 231     INT32: make_integer_parser(4),
 232     INT64: make_integer_parser(8),
 233     DOUBLE: binary64_parser,
 234     BINARY: make_string_parser(lambda b : b),
 235     UTF8: make_string_parser(lambda b : b.decode('utf-8')),
 236     UTF16: make_string_parser(lambda b : b.decode('utf-16')),
 237     UTF32: make_string_parser(lambda b : b.decode('utf-32')),
 238     LIST: list_parser,
 239     DICTIONARY: dictionary_parser,
 240 }
 241
 242 def _object_parser(source):
 243     return _TAGS_TO_PARSERS[source[0]](source[1:])
 244
 245 def _parse(parser, source, consume_all = True):
 246     result = parser(source)
 247
 248     if result.success and result.remaining == b'':
 249         return result.value
 250
 251     raise Exception('Unparsed trailing bytes: {}'.format(result.remaining))
 252
 253 def _binary_deserialize(b):
 254     return _parse(_object_parser, b)
 255
 256 Serializer = collections.namedtuple('Serializer', ['serialize', 'deserialize'])
 257
 258 binary = Serializer(
 259     serialize = _binary_serialize,
 260     deserialize = _binary_deserialize,
 261 )
 262
 263 def binary_to_string(b):
 264     return string.serialize(binary.deserialize(b))
 265
 266 def string_to_binary(s):
 267     return binary.serialize(string.deserialize(s))