1
0
mirror of https://github.com/trezor/trezor-firmware.git synced 2024-11-22 15:38:11 +00:00
trezor-firmware/tools/pb2py
2018-02-27 00:12:44 +01:00

241 lines
7.3 KiB
Python
Executable File

#!/usr/bin/env python3
# Converts Google's protobuf python definitions of TREZOR wire messages
# to plain-python objects as used in TREZOR Core and python-trezor
import sys
import os
import argparse
from google.protobuf.internal.enum_type_wrapper import EnumTypeWrapper
def process_type(t, cls, msg_id, indexfile, is_upy):
print(" * type %s" % t)
imports = []
out = ["", "", "class %s(p.MessageType):" % t, ]
args = []
assigns = []
if cls.DESCRIPTOR.fields_by_name:
out.append(" FIELDS = {")
elif msg_id is None:
out.append(" pass")
for v in sorted(cls.DESCRIPTOR.fields_by_name.values(), key=lambda x: x.number):
number = v.number
fieldname = v.name
type = None
repeated = v.label == 3
required = v.label == 2
# print v.has_default_value, v.default_value
if v.type in (4, 13, 14):
# TYPE_UINT64 = 4
# TYPE_UINT32 = 13
# TYPE_ENUM = 14
type = 'p.UVarintType'
pytype = 'int'
elif v.type in (17,):
# TYPE_SINT32 = 17
type = 'p.Sint32Type'
pytype = 'int'
elif v.type in (18,):
# TYPE_SINT64 = 18
type = 'p.Sint64Type'
pytype = 'int'
elif v.type == 9:
# TYPE_STRING = 9
type = 'p.UnicodeType'
pytype = 'str'
elif v.type == 8:
# TYPE_BOOL = 8
type = 'p.BoolType'
pytype = 'bool'
elif v.type == 12:
# TYPE_BYTES = 12
type = 'p.BytesType'
pytype = 'bytes'
elif v.type == 11:
# TYPE_MESSAGE = 1
type = v.message_type.name
pytype = v.message_type.name
imports.append("from .%s import %s" %
(v.message_type.name, v.message_type.name))
else:
raise Exception("Unknown field type %s for field %s" %
(v.type, fieldname))
if required:
comment = ' # required'
elif v.has_default_value:
comment = ' # default=%s' % repr(v.default_value)
else:
comment = ''
if repeated:
flags = 'p.FLAG_REPEATED'
pytype = "list"
value = None
else:
flags = '0'
value = None
out.append(" %d: ('%s', %s, %s),%s" %
(number, fieldname, type, flags, comment))
args.append(" %s: %s = %s," % (fieldname, pytype, value))
if repeated:
assigns.append(" self.%s = [] if %s is None else %s" % (fieldname, fieldname, fieldname))
else:
assigns.append(" self.%s = %s" % (fieldname, fieldname))
# print fieldname, number, type, repeated, comment
# print v.__dict__
# print v.CPPTYPE_STRING
# print v.LABEL_REPEATED
# print v.enum_type
# v.has_default_value, v.default_value
# v.label == 3 # repeated
# print v.number
if cls.DESCRIPTOR.fields_by_name:
out.append(" }")
if msg_id is not None:
out.append(" MESSAGE_WIRE_TYPE = %d" % msg_id)
if indexfile is not None:
if is_upy:
indexfile.write("%s = const(%d)\n" % (t, msg_id))
else:
indexfile.write("%s = %d\n" % (t, msg_id))
# Remove duplicate imports
imports = sorted(list(set(imports)))
if is_upy:
imports = ['import protobuf as p'] + imports
else:
imports = ['from __future__ import absolute_import',
'from .. import protobuf as p'] + imports
args.append(" **kwargs,")
assigns.append(" p.MessageType.__init__(self, **kwargs)")
init = ["", " def __init__(", " self,"] + args + [" ):"] + assigns
return imports + out + init
def process_enum(t, cls, is_upy):
out = []
if is_upy:
out += ("from micropython import const", "")
print(" * enum %s" % t)
for k, v in cls.items():
# Remove type name from the beginning of the constant
# For example "PinMatrixRequestType_Current" -> "Current"
if k.startswith("%s_" % t):
k = k.replace("%s_" % t, '')
# If type ends with *Type, but constant use type name without *Type, remove it too :)
# For example "ButtonRequestType & ButtonRequest_Other" => "Other"
if t.endswith("Type") and k.startswith("%s_" % t.replace("Type", '')):
k = k.replace("%s_" % t.replace("Type", ''), '')
if is_upy:
out.append("%s = const(%s)" % (k, v))
else:
out.append("%s = %s" % (k, v))
return out
def find_msg_type(msg_types, t):
for k, v in msg_types:
msg_name = k.replace('MessageType_', '')
if msg_name == t:
return v
def process_module(mod, genpath, indexfile, modlist, is_upy):
print("Processing module %s" % mod.__name__)
types = dict([(name, cls)
for name, cls in mod.__dict__.items() if isinstance(cls, type)])
msg_types = __import__('pb2', globals(), locals(), [
'messages_pb2', ]).messages_pb2.MessageType.items()
for t, cls in sorted(types.items()):
# Find message type for given class
msg_id = find_msg_type(msg_types, t)
out = process_type(t, cls, msg_id, indexfile, is_upy)
write_to_file(genpath, t, out)
if modlist:
modlist.write("from .%s import *\n" % t)
enums = dict([(name, cls) for name, cls in mod.__dict__.items()
if isinstance(cls, EnumTypeWrapper)])
for t, cls in enums.items():
out = process_enum(t, cls, is_upy)
write_to_file(genpath, t, out)
if modlist:
modlist.write("from . import %s\n" % t)
def write_to_file(genpath, t, out):
# Write generated sourcecode to given file
f = open(os.path.join(genpath, "%s.py" % t), 'w')
out = ["# Automatically generated by pb2py"] + out
data = "\n".join(out) + "\n"
f.write(data)
f.close()
if __name__ == '__main__':
parser = argparse.ArgumentParser()
parser.add_argument('modulename', type=str, help="Name of module to generate")
parser.add_argument('genpath', type=str, help="Directory for generated source code")
parser.add_argument('-i', '--indexfile', type=str, help="[optional] Generate index file of wire types")
parser.add_argument('-l', '--modlist', type=str, help="[optional] Generate list of modules")
parser.add_argument('-p', '--protopath', type=str, help="[optional] Path to search for pregenerated Google's python sources")
parser.add_argument('-m', '--micropython', action='store_true', help="Use micropython-favoured source code")
args = parser.parse_args()
if args.indexfile:
indexfile = open(args.indexfile, 'a')
else:
indexfile = None
if args.modlist:
modlist = open(args.modlist, 'a')
else:
modlist = None
if args.protopath:
sys.path.append(args.protopath)
# Dynamically load module from argv[1]
tmp = __import__('pb2', globals(), locals(), ['%s_pb2' % args.modulename])
mod = getattr(tmp, "%s_pb2" % args.modulename)
process_module(mod, args.genpath, indexfile, modlist, args.micropython)