Allow pgen to produce a DOT format dump of the grammar (GH-18005)

Originally suggested by Anthony Shaw.

3 files changed, 46 insertions, 2 deletions

diff --git a/Parser/pgen/__main__.py b/Parser/pgen/__main__.py
index bb96e75..d3780a7 100644
--- a/Parser/pgen/__main__.py
+++ b/Parser/pgen/__main__.py
@@ -21,9 +21,19 @@ def main():
     )
 
     parser.add_argument("--verbose", "-v", action="count")
+    parser.add_argument(
+        "--graph",
+        type=argparse.FileType("w"),
+        action="store",
+        metavar="GRAPH_OUTPUT_FILE",
+        help="Dumps a DOT representation of the generated automata in a file",
+    )
+
     args = parser.parse_args()
 
-    p = ParserGenerator(args.grammar, args.tokens, verbose=args.verbose)
+    p = ParserGenerator(
+        args.grammar, args.tokens, verbose=args.verbose, graph_file=args.graph
+    )
     grammar = p.make_grammar()
     grammar.produce_graminit_h(args.graminit_h.write)
     grammar.produce_graminit_c(args.graminit_c.write)
diff --git a/Parser/pgen/automata.py b/Parser/pgen/automata.py
index 545a737..d04ca7c 100644
--- a/Parser/pgen/automata.py
+++ b/Parser/pgen/automata.py
@@ -48,6 +48,26 @@ class NFA:
                 else:
                     writer("    %s -> %d" % (label, j))
 
+    def dump_graph(self, writer):
+        """Dump a DOT representation of the NFA"""
+        writer('digraph %s_nfa {\n' % self.name)
+        todo = [self.start]
+        for i, state in enumerate(todo):
+            writer(' %d [label="State %d %s"];\n' % (i, i, state is self.end and "(final)" or ""))
+            for arc in state.arcs:
+                label = arc.label
+                next = arc.target
+                if next in todo:
+                    j = todo.index(next)
+                else:
+                    j = len(todo)
+                    todo.append(next)
+                if label is None:
+                    writer(" %d -> %d [style=dotted label=ε];\n" % (i, j))
+                else:
+                    writer(" %d -> %d [label=%s];\n" % (i, j, label.replace("'", '"')))
+        writer('}\n')
+
 
 class NFAArc:
     """An arc representing a transition between two NFA states.
@@ -301,6 +321,15 @@ class DFA:
             for label, next in sorted(state.arcs.items()):
                 writer("    %s -> %d" % (label, self.states.index(next)))
 
+    def dump_graph(self, writer):
+        """Dump a DOT representation of the DFA"""
+        writer('digraph %s_dfa {\n' % self.name)
+        for i, state in enumerate(self.states):
+            writer(' %d [label="State %d %s"];\n' % (i, i, state.is_final and "(final)" or ""))
+            for label, next in sorted(state.arcs.items()):
+                writer(" %d -> %d [label=%s];\n" % (i, self.states.index(next), label.replace("'", '"')))
+        writer('}\n')
+
 
 class DFAState(object):
     """A state of a DFA
diff --git a/Parser/pgen/pgen.py b/Parser/pgen/pgen.py
index 2f444eb..03032d4 100644
--- a/Parser/pgen/pgen.py
+++ b/Parser/pgen/pgen.py
@@ -130,7 +130,7 @@ class Label(str):
 
 
 class ParserGenerator(object):
-    def __init__(self, grammar_file, token_file, verbose=False):
+    def __init__(self, grammar_file, token_file, verbose=False, graph_file=None):
         with open(grammar_file) as f:
             self.grammar = f.read()
         with open(token_file) as tok_file:
@@ -141,6 +141,7 @@ class ParserGenerator(object):
         self.opmap["<>"] = "NOTEQUAL"
         self.verbose = verbose
         self.filename = grammar_file
+        self.graph_file = graph_file
         self.dfas, self.startsymbol = self.create_dfas()
         self.first = {}  # map from symbol name to set of tokens
         self.calculate_first_sets()
@@ -152,11 +153,15 @@ class ParserGenerator(object):
             if self.verbose:
                 print("Dump of NFA for", nfa.name)
                 nfa.dump()
+            if self.graph_file is not None:
+                nfa.dump_graph(self.graph_file.write)
             dfa = DFA.from_nfa(nfa)
             if self.verbose:
                 print("Dump of DFA for", dfa.name)
                 dfa.dump()
             dfa.simplify()
+            if self.graph_file is not None:
+                dfa.dump_graph(self.graph_file.write)
             rule_to_dfas[dfa.name] = dfa
 
             if start_nonterminal is None: