1
0
mirror of https://github.com/trezor/trezor-firmware.git synced 2024-12-23 14:58:09 +00:00
trezor-firmware/trezorlib/mapping.py

34 lines
883 B
Python

import trezor_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()