1 files changed, 71 insertions, 0 deletions
diff --git a/lib/protocol_benchmarks.py b/lib/protocol_benchmarks.py
index 81ff365..d2f6188 100755
--- a/lib/protocol_benchmarks.py
+++ b/lib/protocol_benchmarks.py
@@ -2,6 +2,10 @@ import avro.schema
 from avro.datafile import DataFileWriter
 from avro.io import DatumWriter
 
+import thriftpy
+from thriftpy.protocol import TCyBinaryProtocolFactory
+from thriftpy.transport import TCyMemoryBuffer
+
 import bson
 import cbor
 import json
@@ -157,6 +161,70 @@ class Avro(DummyProtocol):
                 self.add_to_dict(new_field['type']['fields'], key, value)
         fields.append(new_field)
 
+class Thrift(DummyProtocol):
+
+    class_index = 1
+
+    def __init__(self, data):
+        super().__init__()
+        self.data = data
+        self._field_id = 1
+        self.proto_buf = ''
+        self.proto_from_json(data)
+
+        with open('/tmp/test.thrift', 'w') as f:
+            f.write(self.proto_buf)
+
+        membuf = TCyMemoryBuffer()
+        proto = TCyBinaryProtocolFactory().get_protocol(membuf)
+        # TODO irgendwo bleibt state übrig -> nur das bei allerersten
+        # Aufruf geladene Protokoll wird berücksichtigt, dazu nicht passende
+        # Daten werden nicht serialisiert
+        test_thrift = thriftpy.load('/tmp/test.thrift', module_name='test{:d}_thrift'.format(Thrift.class_index))
+        Thrift.class_index += 1
+        benchmark = test_thrift.Benchmark()
+
+        for key, value in data.items():
+            benchmark.__dict__[key] = value
+
+        try:
+            proto.write_struct(benchmark)
+        except thriftpy.thrift.TDecodeException:
+            raise RuntimeError('Unsupported data layout') from None
+        membuf.flush()
+        self.serialized_data = membuf.getvalue()
+
+    def can_get_serialized_length(self):
+        return True
+
+    def get_serialized_length(self):
+        return len(self.serialized_data)
+
+    def type_to_type_name(self, value):
+        type_type = type(value)
+        if type_type == int:
+            return 'i32'
+        if type_type == float:
+            return 'double'
+        if type_type == str:
+            return 'string'
+        if type_type == list:
+            return 'list<{}>'.format(self.type_to_type_name(value[0]))
+        if type_type == dict:
+            sub_value = list(value.values())[0]
+            return 'map<{},{}>'.format('string', self.type_to_type_name(sub_value))
+
+    def add_to_dict(self, key, value):
+        key_type = self.type_to_type_name(value)
+        self.proto_buf += '{:d}: {} {};\n'.format(self._field_id, key_type, key)
+        self._field_id += 1
+
+    def proto_from_json(self, data):
+        self.proto_buf += 'struct Benchmark {\n'
+        for key, value in data.items():
+            self.add_to_dict(key, value)
+        self.proto_buf += '}\n'
+
 class ArduinoJSON(DummyProtocol):
 
     def __init__(self, data, bufsize = 255, int_type = 'uint16_t', float_type = 'float'):
@@ -1403,6 +1471,9 @@ def codegen_for_lib(library, library_options, data):
             strbuf = int(strbuf)
         return NanoPB(data, cardinality = cardinality, max_string_length = strbuf)
 
+    if library == 'thrift':
+        return Thrift(data)
+
     if library == 'ubjson':
         return UBJ(data)