protob: refactor graph.py

protob/graph.py

@@ -7,12 +7,24 @@ from graphviz import Digraph
 class Message(object):
     def __init__(self, name, attrs):
         self.name = name
-        self.attrs = attrs
+        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("wrong message type in message '%s'" % m.name)
+        self.next = []
+        for a in attrs:
+            if a[0] == "next":
+                self.next.append(a[1])
 
 
 def generate_messages(files):
     attrs = []
-    msgs = []
+    msgs = {}
     for f in files:
         for line in open(f, "rt").readlines():
             line = line.rstrip()
@@ -20,7 +32,7 @@ def generate_messages(files):
                 attrs.append(line[4:].split(" "))
             elif line.startswith("message "):
                 name = line[8:-2]
-                msgs.append(Message(name, attrs))
+                msgs[name] = Message(name, attrs)
                 attrs = []
     return msgs
 
@@ -28,26 +40,21 @@ def generate_messages(files):
 def generate_graph(msgs, fn):
     dot = Digraph(format="png")
     dot.attr(rankdir="LR")
-    for m in msgs:
-        typ = m.attrs[0][0]
-        if typ == "start":
+    for m in msgs.values():
+        if m.typ == "start":
             dot.node(m.name, shape="box", color="blue")
-        elif typ == "end":
+        elif m.typ == "end":
             dot.node(m.name, shape="box", color="green3")
-        elif typ == "auxstart":
+        elif m.typ == "auxstart":
             dot.node(m.name, shape="diamond", color="blue")
-        elif typ == "auxend":
+        elif m.typ == "auxend":
             dot.node(m.name, shape="diamond", color="green3")
-        elif typ == "next":
-            dot.node(m.name)  # no attrs
-        elif typ in ["embed", "ignore"]:
-            continue
-        else:
-            raise ValueError("wrong message type in message '%s'" % m.name)
-    for m in msgs:
-        for a in m.attrs:
-            if a[0] == "next":
-                dot.edge(m.name, a[1])
+        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)