Add AST to ASCII converter

jasminesnake/ast/__init__.py

@@ -1,7 +1,8 @@
 """AST module."""
 
+from enum import Enum
+from typing import Union
 from antlr4 import ParseTreeWalker
-from tree_format import format_tree
 
 import lex.JavaScriptParser as Parser
 import ast.nodes
@@ -24,6 +25,46 @@ def from_parse_tree(tree: JSP.ProgramContext) -> ast.nodes.Program:
     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

jasminesnake/ast/nodes.py

@@ -21,8 +21,9 @@ Todo:
     * Add support for lacking features
 """
 
-from typing import List, Union, Optional, Literal as TypeLiteral, TypedDict
+from typing import List, Union, Optional, Literal as TypeLiteral, TypedDict, Any
 from enum import Enum
+from collections import OrderedDict
 
 # The Lord sees I actually wanted to split it up, but ESTree hierarchy is so messed up... No. It's actually *fucked up*
 # that much that I couldn't even resolve circular dependencies in the submodules. I have to reap what I've sown.
@@ -140,6 +141,10 @@ class Property:
     ...
 
 
+class Identifier:
+    ...
+
+
 # "Node objects" block
 
 
@@ -155,6 +160,9 @@ class Position:
         self.line = line
         self.column = column
 
+    def __str__(self):
+        return f"{self.line}:{self.column}"
+
 
 class SourceLocation:
     """
@@ -172,6 +180,10 @@ class SourceLocation:
         self.start = start
         self.end = end
 
+    def __str__(self):
+        src = "" if self.source is None else f"{self.source}:"
+        return f"{src}{str(self.start)}"
+
 
 class Node:
     """ESTree AST nodes are represented as Node objects, which may have any prototype inheritance but which implement
@@ -191,16 +203,15 @@ class Node:
         self.type = node_type
         self.loc = loc
 
+        self._fields: OrderedDict[str, Any] = OrderedDict()
+        self._fields.update({"type": self.type, "loc": self.loc})
 
-# "Identifier" block
+    def __str__(self):
+        return f"{self.type} at {str(self.loc)}"
 
-
+    @property
-class Identifier(Expression, Pattern):
+    def fields(self):
-    """An identifier. Note that an identifier may be an expression or a destructuring pattern."""
+        return self._fields
-
-    def __init__(self, loc: Optional[SourceLocation], name: str):
-        super(Identifier, self).__init__("Identifier", loc)
-        self.name = name
 
 
 # "Literal" block
@@ -214,6 +225,7 @@ class Literal(Expression):
     ):
         super().__init__("Literal", loc)
         self.value = value
+        self._fields.update({"value": self.value})
 
 
 # "Programs" block
@@ -231,6 +243,7 @@ class Program(Node):
         super().__init__("Program", loc)
         self.body = body
         self.source_type = source_type
+        self._fields.update({"sourceType": self.source_type, "body": self.body})
 
 
 # "Functions" block
@@ -257,6 +270,7 @@ class Function(Node):
         self.id = function_id
         self.params = params
         self.body = body
+        self._fields.update({"id": self.id, "params": self.params, "body": self.body})
 
 
 # "Statements" block
@@ -282,6 +296,7 @@ class BlockStatement(Statement):
     def __init__(self, loc: Optional[SourceLocation], body: List[Statement]):
         super().__init__("BlockStatement", loc)
         self.body = body
+        self._fields.update({"body": self.body})
 
 
 class ExpressionStatement(Statement):
@@ -290,6 +305,7 @@ class ExpressionStatement(Statement):
     def __init__(self, loc: Optional[SourceLocation], expression: Expression):
         super().__init__("ExpressionStatement", loc)
         self.expression = expression
+        self._fields.update({"expression": self.expression})
 
 
 class Directive(Node):
@@ -303,6 +319,9 @@ class Directive(Node):
         super().__init__("Directive", loc)
         self.expression = expression
         self.directive = directive
+        self._fields.update(
+            {"expression": self.expression, "directive": self.directive}
+        )
 
 
 class FunctionBody(BlockStatement):
@@ -320,6 +339,7 @@ class ReturnStatement(Statement):
     def __init__(self, loc: Optional[SourceLocation], argument: Optional[Expression]):
         super().__init__("ReturnStatement", loc)
         self.argument = argument
+        self._fields.update({"argument": self.argument})
 
 
 class BreakStatement(Statement):
@@ -328,6 +348,7 @@ class BreakStatement(Statement):
     def __init__(self, loc: Optional[SourceLocation], label: Optional[Identifier]):
         super().__init__("BreakStatement", loc)
         self.label = label
+        self._fields.update({"label": self.label})
 
 
 class ContinueStatement(Statement):
@@ -336,6 +357,7 @@ class ContinueStatement(Statement):
     def __init__(self, loc: Optional[SourceLocation], label: Optional[Identifier]):
         super().__init__("ContinueStatement", loc)
         self.label = label
+        self._fields.update({"label": self.label})
 
 
 class IfStatement(Statement):
@@ -352,6 +374,13 @@ class IfStatement(Statement):
         self.test = test
         self.consequent = consequent
         self.alternate = alternate
+        self._fields.update(
+            {
+                "test": self.test,
+                "consequent": self.consequent,
+                "alternate": self.alternate,
+            }
+        )
 
 
 class WhileStatement(Statement):
@@ -363,6 +392,7 @@ class WhileStatement(Statement):
         super().__init__("WhileStatement", loc)
         self.test = test
         self.body = body
+        self._fields.update({"test": self.test, "body": self.body})
 
 
 class DoWhileStatement(Statement):
@@ -374,6 +404,7 @@ class DoWhileStatement(Statement):
         super().__init__("DoWhileStatement", loc)
         self.body = body
         self.test = test
+        self._fields.update({"body": self.body, "test": self.test})
 
 
 class ForStatement(Statement):
@@ -392,6 +423,14 @@ class ForStatement(Statement):
         self.test = test
         self.update = update
         self.body = body
+        self._fields.update(
+            {
+                "init": self.init,
+                "test": self.test,
+                "update": self.update,
+                "body": self.body,
+            }
+        )
 
 
 class ForInStatement(Statement):
@@ -408,6 +447,7 @@ class ForInStatement(Statement):
         self.left = left
         self.right = right
         self.body = body
+        self._fields.update({"left": self.left, "right": self.right, "body": self.body})
 
 
 # "Declarations" block
@@ -443,6 +483,7 @@ class VariableDeclarator(Node):
         super().__init__("VariableDeclarator", loc)
         self.id = var_id
         self.init = init
+        self._fields.update({"id": self.id, "init": self.init})
 
 
 class VariableDeclaration(Declaration):
@@ -457,6 +498,7 @@ class VariableDeclaration(Declaration):
         super().__init__("VariableDeclaration", loc)
         self.declarations = declarations
         self.kind = kind
+        self._fields.update({"kind": self.kind, "declarations": self.declarations})
 
 
 # "Expressions" block
@@ -487,6 +529,7 @@ class SpreadElement(Node):
     def __init__(self, loc: Optional[SourceLocation], argument: Expression):
         super().__init__("SpreadElement", loc)
         self.argument = argument
+        self._fields.update({"argument": self.argument})
 
 
 class ThisExpression(Expression):
@@ -506,6 +549,7 @@ class ArrayExpression(Expression):
     ):
         super().__init__("ArrayExpression", loc)
         self.elements = elements
+        self._fields.update({"elements": self.elements})
 
 
 class ObjectExpression(Expression):
@@ -514,6 +558,7 @@ class ObjectExpression(Expression):
     def __init__(self, loc: Optional[SourceLocation], properties: List[Property]):
         super().__init__("ObjectExpression", loc)
         self.properties = properties
+        self._fields.update({"properties": self.properties})
 
 
 class FunctionExpression(Function, Expression):
@@ -541,6 +586,7 @@ class ArrowFunctionExpression(Function, Expression):
     ):
         super().__init__("ArrowFunctionExpression", loc, None, params, body)
         self.expression = expression
+        self._fields.update({"expression": self.expression})
 
 
 class UnaryExpression(Expression):
@@ -557,6 +603,13 @@ class UnaryExpression(Expression):
         self.operator = operator
         self.prefix = prefix
         self.argument = argument
+        self._fields.update(
+            {
+                "operator": self.operator,
+                "prefix": self.prefix,
+                "argument": self.argument,
+            }
+        )
 
 
 class UpdateExpression(Expression):
@@ -573,6 +626,13 @@ class UpdateExpression(Expression):
         self.operator = operator
         self.argument = argument
         self.prefix = prefix
+        self._fields.update(
+            {
+                "operator": self.operator,
+                "argument": self.argument,
+                "prefix": self.prefix,
+            }
+        )
 
 
 class BinaryExpression(Expression):
@@ -589,6 +649,9 @@ class BinaryExpression(Expression):
         self.operator = operator
         self.left = left
         self.right = right
+        self._fields.update(
+            {"operator": self.operator, "left": self.left, "right": self.right}
+        )
 
 
 class AssignmentExpression(Expression):
@@ -607,6 +670,9 @@ class AssignmentExpression(Expression):
         self.operator = operator
         self.left = left
         self.right = right
+        self._fields.update(
+            {"operator": self.operator, "left": self.left, "right": self.right}
+        )
 
 
 class LogicalExpression(Expression):
@@ -623,6 +689,9 @@ class LogicalExpression(Expression):
         self.operator = operator
         self.left = left
         self.right = right
+        self._fields.update(
+            {"operator": self.operator, "left": self.left, "right": self.right}
+        )
 
 
 class MemberExpression(Expression, Pattern):
@@ -641,6 +710,13 @@ class MemberExpression(Expression, Pattern):
         self.object = member_object
         self.property = member_property
         self.computed = computed
+        self._fields.update(
+            {
+                "object": self.object,
+                "property": self.property,
+                "computed": self.computed,
+            }
+        )
 
 
 class ConditionalExpression(Expression):
@@ -657,6 +733,13 @@ class ConditionalExpression(Expression):
         self.test = test
         self.alternate = alternate
         self.consequent = consequent
+        self._fields.update(
+            {
+                "test": self.test,
+                "alternate": self.alternate,
+                "consequent": self.consequent,
+            }
+        )
 
 
 class CallExpression(Expression):
@@ -671,6 +754,7 @@ class CallExpression(Expression):
         super().__init__("CallExpression", loc)
         self.callee = callee
         self.arguments = arguments
+        self._fields.update({"callee": self.callee, "arguments": self.arguments})
 
 
 class NewExpression(Expression):
@@ -685,6 +769,7 @@ class NewExpression(Expression):
         super().__init__("NewExpression", loc)
         self.callee = callee
         self.arguments = arguments
+        self._fields.update({"callee": self.callee, "arguments": self.arguments})
 
 
 class SequenceExpression(Expression):
@@ -693,6 +778,7 @@ class SequenceExpression(Expression):
     def __init__(self, loc: Optional[SourceLocation], expressions: List[Expression]):
         super().__init__("SequenceExpression", loc)
         self.expressions = expressions
+        self._fields.update({"expressions": self.expressions})
 
 
 def _generate_unary_expression(operator: UnaryOperator, docstring: str):
@@ -940,6 +1026,16 @@ class Property(Node):
         self.method = method
         self.shorthand = shorthand
         self.computed = computed
+        self._fields.update(
+            {
+                "key": self.key,
+                "value": self.value,
+                "kind": self.kind,
+                "method": self.method,
+                "shorthand": self.shorthand,
+                "computed": self.computed,
+            }
+        )
 
 
 class AssignmentProperty(Property):
@@ -978,6 +1074,7 @@ class ObjectPattern(Pattern):
     ):
         super().__init__("ObjectPattern", loc)
         self.properties = properties
+        self._fields.update({"properties": self.properties})
 
 
 class ArrayPattern(Pattern):
@@ -986,3 +1083,15 @@ class ArrayPattern(Pattern):
     ):
         super().__init__("ArrayPattern", loc)
         self.elements = elements
+        self._fields.update({"elements": self.elements})
+
+
+# "Identifier" block
+
+
+class Identifier(Expression, Pattern):
+    """An identifier. Note that an identifier may be an expression or a destructuring pattern."""
+
+    def __init__(self, loc: Optional[SourceLocation], name: str):
+        super().__init__("Identifier", loc)
+        self.name = name

requirements.txt

@@ -1,4 +1,3 @@
 antlr4-python3-runtime==4.8
 colorama==0.4.3
-coloredlogs==14.0
+coloredlogs==14.0
-tree_format==0.1.2