1import argparse
2import sys
3from typing import Any, Callable, Iterator
4
5from pegen.build import build_parser
6from pegen.grammar import Grammar, Rule
7
8argparser = argparse.ArgumentParser(
9    prog="pegen", description="Pretty print the AST for a given PEG grammar"
10)
11argparser.add_argument("filename", help="Grammar description")
12
13
14class ASTGrammarPrinter:
15    def children(self, node: Rule) -> Iterator[Any]:
16        for value in node:
17            if isinstance(value, list):
18                yield from value
19            else:
20                yield value
21
22    def name(self, node: Rule) -> str:
23        if not list(self.children(node)):
24            return repr(node)
25        return node.__class__.__name__
26
27    def print_grammar_ast(self, grammar: Grammar, printer: Callable[..., None] = print) -> None:
28        for rule in grammar.rules.values():
29            printer(self.print_nodes_recursively(rule))
30
31    def print_nodes_recursively(self, node: Rule, prefix: str = "", istail: bool = True) -> str:
32
33        children = list(self.children(node))
34        value = self.name(node)
35
36        line = prefix + ("└──" if istail else "├──") + value + "\n"
37        sufix = "   " if istail else "│  "
38
39        if not children:
40            return line
41
42        *children, last = children
43        for child in children:
44            line += self.print_nodes_recursively(child, prefix + sufix, False)
45        line += self.print_nodes_recursively(last, prefix + sufix, True)
46
47        return line
48
49
50def main() -> None:
51    args = argparser.parse_args()
52
53    try:
54        grammar, parser, tokenizer = build_parser(args.filename)
55    except Exception as err:
56        print("ERROR: Failed to parse grammar file", file=sys.stderr)
57        sys.exit(1)
58
59    visitor = ASTGrammarPrinter()
60    visitor.print_grammar_ast(grammar)
61
62
63if __name__ == "__main__":
64    main()
65