mirror of https://github.com/t1meshift/js.git
71 lines
1.7 KiB
Python
71 lines
1.7 KiB
Python
"""AST module."""
|
|
|
|
from enum import Enum
|
|
from typing import Union
|
|
from antlr4 import ParseTreeWalker
|
|
|
|
import lex.JavaScriptParser as Parser
|
|
import ast.nodes
|
|
from .parse_tree_listeners import ASTListener
|
|
|
|
JSP = Parser.JavaScriptParser
|
|
|
|
|
|
def from_parse_tree(tree: JSP.ProgramContext) -> ast.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
|
|
|
|
|
|
def to_ascii_tree(
|
|
node: Union[ast.nodes.Position, ast.nodes.SourceLocation, ast.nodes.Node],
|
|
name_prefix: str = "",
|
|
nesting_lvl: int = 0,
|
|
):
|
|
if nesting_lvl < 0:
|
|
raise ValueError("Nesting level can't be below 0")
|
|
|
|
FORK = "+"
|
|
VERTICAL = "|"
|
|
HORIZONTAL = "-"
|
|
|
|
SUBENTRY_PREFIX = f"{FORK}{HORIZONTAL}{HORIZONTAL} "
|
|
NESTED_PREFIX = f"{VERTICAL} "
|
|
|
|
value = str(node)
|
|
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"
|
|
|
|
if children is not None:
|
|
for (child_name, child_value) in children:
|
|
result += to_ascii_tree(child_value, f"{child_name}: ", nesting_lvl + 1)
|
|
|
|
# result += "\n"
|
|
|
|
return result
|
|
|
|
|
|
# Delete temporary imports
|
|
del JSP
|
|
del Parser
|