Add short mode for AST printing

master
Yury Kurlykov 2020-06-28 03:05:00 +10:00
parent 5cd616ceb3
commit 5024ca003b
Signed by: t1meshift
GPG Key ID: B133F3167ABF94D8
2 changed files with 20 additions and 6 deletions

View File

@ -20,6 +20,9 @@ def create_argument_parser():
) )
_arg_parser.add_argument("--snake", action="store_true", help="print a snake") _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( _arg_parser.add_argument(
"--verbose", "--verbose",
"-v", "-v",
@ -72,9 +75,12 @@ def main():
tree = stream.parse() tree = stream.parse()
ast_tree = from_parse_tree(tree) 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 # TODO: run logic
sys.exit(0) sys.exit(0)

View File

@ -29,10 +29,13 @@ def to_ascii_tree(
node: Union[nodes.Position, nodes.SourceLocation, nodes.Node], node: Union[nodes.Position, nodes.SourceLocation, nodes.Node],
name_prefix: str = "", name_prefix: str = "",
nesting_lvl: int = 0, nesting_lvl: int = 0,
ast_format: str = "full",
): ):
if nesting_lvl < 0: if nesting_lvl < 0:
raise ValueError("Nesting level can't be below 0") raise ValueError("Nesting level can't be below 0")
NODE_BLACKLIST = ["loc", "type"] if ast_format == "short" else []
FORK = "+" FORK = "+"
VERTICAL = "|" VERTICAL = "|"
HORIZONTAL = "-" HORIZONTAL = "-"
@ -40,25 +43,30 @@ def to_ascii_tree(
SUBENTRY_PREFIX = f"{FORK}{HORIZONTAL}{HORIZONTAL} " SUBENTRY_PREFIX = f"{FORK}{HORIZONTAL}{HORIZONTAL} "
NESTED_PREFIX = f"{VERTICAL} " NESTED_PREFIX = f"{VERTICAL} "
value = str(node) value = str(node) if not isinstance(node, list) else ""
children = None children = None
if isinstance(node, Enum): if isinstance(node, Enum):
value = str(node.value) value = str(node.value)
if isinstance(node, list): if isinstance(node, list):
value = ""
children = [(index, val) for index, val in enumerate(node)] children = [(index, val) for index, val in enumerate(node)]
if hasattr(node, "fields"): if hasattr(node, "fields"):
children = [(k, node.fields[k]) for k in node.fields.keys()] 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"{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: if children is not None:
for (child_name, child_value) in children: 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" # result += "\n"