python-trezor/trezorlib/mapping.py

34 lines
885 B
Python
Raw Normal View History

2013-12-16 07:26:40 -08:00
import messages_pb2 as proto
2012-11-13 06:09:39 -08:00
map_type_to_class = {}
2012-11-13 06:09:39 -08:00
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
2012-11-13 06:09:39 -08:00
def get_type(msg):
return map_class_to_type[msg.__class__]
2012-11-13 06:09:39 -08:00
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)]
2012-11-13 06:09:39 -08:00
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()
2012-11-13 06:09:39 -08:00
check_missing()