improve ast

Author: Scony

gdtoolkit/common/ast.py

@@ -19,10 +19,10 @@ def __init__(self, node: Tree):
 
     def _load_sub_statements(self):
         if self.kind == "class_def":
-            pass
-        elif self.kind == "property_body_def":
-            pass
-        elif self.kind in ["func_def", "static_func_def"]:
+            raise NotImplementedError
+        if self.kind == "property_body_def":
+            raise NotImplementedError
+        if self.kind in ["func_def", "static_func_def"]:
             self.sub_statements = [Statement(n) for n in self.lark_node.children[1:]]
         elif self.kind == "if_stmt":
             for branch in self.lark_node.children:
@@ -31,16 +31,22 @@ def _load_sub_statements(self):
                 else:
                     self.sub_statements += [Statement(n) for n in branch.children]
         elif self.kind == "while_stmt":
-            pass
+            self.sub_statements = [Statement(n) for n in self.lark_node.children[1:]]
         elif self.kind == "for_stmt":
-            pass
+            self.sub_statements = [Statement(n) for n in self.lark_node.children[2:]]
         elif self.kind == "match_stmt":
-            pass
+            for branch in self.lark_node.children:
+                self.sub_statements += [Statement(n) for n in branch.children[1:]]
         for sub_statement in self.sub_statements:
             self.all_sub_statements += [
                 sub_statement
             ] + sub_statement.all_sub_statements
 
+    def __repr__(self):
+        return "Statement({}:{}:{})".format(
+            self.lark_node.data, self.lark_node.line, self.lark_node.column
+        )
+
 
 # pylint: disable=too-few-public-methods
 class Parameter:
@@ -50,15 +56,15 @@ def __init__(self, node: Tree):
         self.name = node.children[0].value
 
 
-# TODO: inherit from statement
 # pylint: disable=too-few-public-methods
-class Function:
+class Function(Statement):
     """Abstract representation of function"""
 
     def __init__(self, func_def: Tree):
-        self.lark_node = func_def
+        super().__init__(func_def)
         self.name = ""
         self.parameters = []  # type: List[Parameter]
+        self.all_statements = [self] + self.all_sub_statements  # type: ignore
 
         self._load_data_from_func_def(func_def)
 
@@ -72,8 +78,6 @@ def _load_data_from_func_def(self, func_def: Tree) -> None:
             for c in func_args.children  # type: ignore
             if c.data != "trailing_comma"
         ]
-        slf = Statement(self.lark_node)
-        self.all_statements = [slf] + slf.all_sub_statements
 
 
 # pylint: disable=too-few-public-methods

tests/test_ast.py

@@ -0,0 +1,59 @@
+from gdtoolkit.parser import parser
+from gdtoolkit.common.ast import AbstractSyntaxTree
+
+
+def test_toplevel():
+    code = """func foo():
+    pass
+    """
+    parse_tree = parser.parse(code, gather_metadata=True)
+    ast = AbstractSyntaxTree(parse_tree)
+    assert len(ast.all_functions) == 1
+    function = ast.all_functions[0]
+    assert function.name == "foo"
+    assert len(function.all_sub_statements) == 1
+
+
+def test_all_sub_statements_of_function():
+    code = """func foo():
+    pass
+    if true:
+        pass
+        if true:
+            return
+        elif true:
+            return
+        else:
+            return
+    while true:
+        if true:
+            return
+    for f in range(10):
+        if true:
+            return
+    var x
+    match(x):
+        1:
+            pass
+    """
+    parse_tree = parser.parse(code, gather_metadata=True)
+    ast = AbstractSyntaxTree(parse_tree)
+    assert len(ast.all_functions) == 1
+    function = ast.all_functions[0]
+    assert len(function.all_sub_statements) == 16
+    assert function.all_sub_statements[0].kind == 'pass_stmt'
+    assert function.all_sub_statements[1].kind == 'if_stmt'
+    assert function.all_sub_statements[2].kind == 'pass_stmt'
+    assert function.all_sub_statements[3].kind == 'if_stmt'
+    assert function.all_sub_statements[4].kind == 'return_stmt'
+    assert function.all_sub_statements[5].kind == 'return_stmt'
+    assert function.all_sub_statements[6].kind == 'return_stmt'
+    assert function.all_sub_statements[7].kind == 'while_stmt'
+    assert function.all_sub_statements[8].kind == 'if_stmt'
+    assert function.all_sub_statements[9].kind == 'return_stmt'
+    assert function.all_sub_statements[10].kind == 'for_stmt'
+    assert function.all_sub_statements[11].kind == 'if_stmt'
+    assert function.all_sub_statements[12].kind == 'return_stmt'
+    assert function.all_sub_statements[13].kind == 'func_var_stmt'
+    assert function.all_sub_statements[14].kind == 'match_stmt'
+    assert function.all_sub_statements[15].kind == 'pass_stmt'