python-trezor/trezorlib/mapping.py
2013-12-16 17:40:25 +01:00

34 lines
885 B
Python

import messages_pb2 as proto
map_type_to_class = {}
map_class_to_type = {}
def build_map():
for msg_type, i in proto.MessageType.items():
msg_name = msg_type.replace('MessageType_', '')
msg_class = getattr(proto, msg_name)
map_type_to_class[i] = msg_class
map_class_to_type[msg_class] = i
def get_type(msg):
return map_class_to_type[msg.__class__]
def get_class(t):
return map_type_to_class[t]
def check_missing():
from google.protobuf import reflection
types = [getattr(proto, item) for item in dir(proto)
if issubclass(getattr(proto, item).__class__, reflection.GeneratedProtocolMessageType)]
missing = list(set(types) - set(map_type_to_class.values()))
if len(missing):
raise Exception("Following protobuf messages are not defined in mapping: %s" % missing)
build_map()
check_missing()