1
0
mirror of https://github.com/trezor/trezor-firmware.git synced 2024-11-22 15:38:11 +00:00
trezor-firmware/common/protob/graph.py
2021-10-27 13:13:15 +02:00

77 lines
2.0 KiB
Python
Executable File

#!/usr/bin/env python3
import sys
from os.path import basename
from graphviz import Digraph
class Message:
def __init__(self, fname, name, attrs):
self.fname = basename(fname)
self.name = name
if len(attrs) == 0:
raise ValueError(f"message {name} has no attributes")
t = attrs[0][0]
if t in ["start", "end", "auxstart", "auxend", "embed", "ignore"]:
self.typ = t
attrs = attrs[1:]
elif t == "next":
self.typ = "normal"
attrs = attrs
else:
raise ValueError(f"wrong message type in message {name}")
self.next = []
for a in attrs:
if a[0] == "next":
self.next.append(a[1])
def __repr__(self):
return f'{self.name}(type={self.typ}, fname="{self.fname}", next={self.next})'
def generate_messages(files):
attrs = []
msgs = {}
for f in files:
for line in open(f, "rt").readlines():
line = line.rstrip()
if line.startswith(" * @"):
attrs.append(line[4:].split(" "))
elif line.startswith("message "):
name = line[8:-2]
msgs[name] = Message(f, name, attrs)
attrs = []
return msgs
def generate_graph(msgs, fn):
dot = Digraph()
dot.attr(rankdir="LR")
for m in msgs.values():
if m.typ == "start":
dot.node(m.name, shape="box", color="blue")
elif m.typ == "end":
dot.node(m.name, shape="box", color="green3")
elif m.typ == "auxstart":
dot.node(m.name, shape="diamond", color="blue")
elif m.typ == "auxend":
dot.node(m.name, shape="diamond", color="green3")
elif m.typ == "normal":
dot.node(m.name)
for m in msgs.values():
for n in m.next:
dot.edge(m.name, n)
dot.render(fn)
def main():
proto_files = sys.argv
msgs = generate_messages(proto_files)
generate_graph(msgs, "proto.gv")
if __name__ == "__main__":
main()