diff --git a/jasminesnake/__main__.py b/jasminesnake/__main__.py index c86e8a7..3e81cf3 100644 --- a/jasminesnake/__main__.py +++ b/jasminesnake/__main__.py @@ -20,6 +20,9 @@ def create_argument_parser(): ) _arg_parser.add_argument("--snake", action="store_true", help="print a snake") + _arg_parser.add_argument( + "--ast", choices=["full", "short", "none"], default="none", help="print AST" + ) _arg_parser.add_argument( "--verbose", "-v", @@ -72,9 +75,12 @@ def main(): tree = stream.parse() ast_tree = from_parse_tree(tree) - ascii_ast = to_ascii_tree(ast_tree) - logging.info("Got an AST!\n%s", ascii_ast) + logging.info("Got an AST!\n") + if args.ast != "none": + ascii_ast = to_ascii_tree(ast_tree, ast_format=args.ast) + print(ascii_ast) + # TODO: run logic sys.exit(0) diff --git a/jasminesnake/ast/__init__.py b/jasminesnake/ast/__init__.py index 2e5e22a..68beb83 100644 --- a/jasminesnake/ast/__init__.py +++ b/jasminesnake/ast/__init__.py @@ -29,10 +29,13 @@ def to_ascii_tree( node: Union[nodes.Position, nodes.SourceLocation, nodes.Node], name_prefix: str = "", nesting_lvl: int = 0, + ast_format: str = "full", ): if nesting_lvl < 0: raise ValueError("Nesting level can't be below 0") + NODE_BLACKLIST = ["loc", "type"] if ast_format == "short" else [] + FORK = "+" VERTICAL = "|" HORIZONTAL = "-" @@ -40,25 +43,30 @@ def to_ascii_tree( SUBENTRY_PREFIX = f"{FORK}{HORIZONTAL}{HORIZONTAL} " NESTED_PREFIX = f"{VERTICAL} " - value = str(node) + value = str(node) if not isinstance(node, list) else "" children = None if isinstance(node, Enum): value = str(node.value) if isinstance(node, list): - value = "" children = [(index, val) for index, val in enumerate(node)] if hasattr(node, "fields"): children = [(k, node.fields[k]) for k in node.fields.keys()] result = f"{NESTED_PREFIX * (nesting_lvl - 1)}{SUBENTRY_PREFIX * (nesting_lvl > 0)}" - result += f"{name_prefix}{value}\n" + result += f"{name_prefix + ' '*(value!='' and name_prefix!='')}{value}\n" if children is not None: for (child_name, child_value) in children: - result += to_ascii_tree(child_value, f"{child_name}: ", nesting_lvl + 1) + if child_name not in NODE_BLACKLIST: + result += to_ascii_tree( + child_value, + f"{child_name}:", + nesting_lvl + 1, + ast_format=ast_format, + ) # result += "\n"