Add short mode for AST printing

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)
 

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"