js/jasminesnake/ast/__init__.py

79 lines
2.0 KiB
Python
Raw Normal View History

2020-04-26 23:10:26 +10:00
"""AST module."""
2020-04-28 03:27:07 +10:00
from enum import Enum
from typing import Union
from antlr4 import ParseTreeWalker
import jasminesnake.lex.JavaScriptParser as Parser
from . import nodes
from .parse_tree_listeners import ASTListener
JSP = Parser.JavaScriptParser
def from_parse_tree(tree: JSP.ProgramContext) -> nodes.Program:
"""Generate AST from ANTLR parse tree.
Args:
tree (JSP.ProgramContext): ANTLR parse tree.
Returns:
`Program` AST node, which is the root node.
"""
ast_listener = ASTListener()
ParseTreeWalker.DEFAULT.walk(ast_listener, tree)
return ast_listener.program_node
2020-04-28 03:27:07 +10:00
def to_ascii_tree(
node: Union[nodes.Position, nodes.SourceLocation, nodes.Node],
2020-04-28 03:27:07 +10:00
name_prefix: str = "",
nesting_lvl: int = 0,
2020-06-28 03:05:00 +10:00
ast_format: str = "full",
2020-04-28 03:27:07 +10:00
):
if nesting_lvl < 0:
raise ValueError("Nesting level can't be below 0")
2020-06-28 03:05:00 +10:00
NODE_BLACKLIST = ["loc", "type"] if ast_format == "short" else []
2020-04-28 03:27:07 +10:00
FORK = "+"
VERTICAL = "|"
HORIZONTAL = "-"
SUBENTRY_PREFIX = f"{FORK}{HORIZONTAL}{HORIZONTAL} "
NESTED_PREFIX = f"{VERTICAL} "
2020-06-28 03:05:00 +10:00
value = str(node) if not isinstance(node, list) else ""
2020-04-28 03:27:07 +10:00
children = None
if isinstance(node, Enum):
value = str(node.value)
if isinstance(node, list):
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)}"
2020-06-28 03:05:00 +10:00
result += f"{name_prefix + ' '*(value!='' and name_prefix!='')}{value}\n"
2020-04-28 03:27:07 +10:00
if children is not None:
for (child_name, child_value) in children:
2020-06-28 03:05:00 +10:00
if child_name not in NODE_BLACKLIST:
result += to_ascii_tree(
child_value,
f"{child_name}:",
nesting_lvl + 1,
ast_format=ast_format,
)
2020-04-28 03:27:07 +10:00
# result += "\n"
return result
# Delete temporary imports
del JSP
del Parser