From cab7a1c129887bb965ded84bdc2e5e627a3febba Mon Sep 17 00:00:00 2001
From: Aleksey Filippov <al.filippov@ulstu.ru>
Date: Fri, 21 Jan 2022 16:16:41 +0400
Subject: [PATCH] Improve parse tree

---
 src/parse_tree/parse_tree.py      | 21 ++++++++++++++++-----
 src/parse_tree/parse_tree_item.py |  3 +++
 2 files changed, 19 insertions(+), 5 deletions(-)

diff --git a/src/parse_tree/parse_tree.py b/src/parse_tree/parse_tree.py
index c7c9d1f..b601f81 100644
--- a/src/parse_tree/parse_tree.py
+++ b/src/parse_tree/parse_tree.py
@@ -1,5 +1,5 @@
 import warnings
-from typing import List
+from typing import List, Dict
 
 from anytree import Node, RenderTree
 
@@ -23,10 +23,21 @@ class ParseTree:
         return nodes
 
     def __create_tree(self, nodes: List[ParseItem]) -> Node:
-        root: Node
-        for node in nodes:
-            if node.parent_index() == 0:
-                root = Node(node)
+        tree_nodes_count: int = 0
+        parents: Dict[int, Node] = {}
+        root = None
+        while tree_nodes_count < len(nodes):
+            for node in nodes:
+                if parents.get(node.index()) is not None:
+                    continue
+                parent: Node = parents.get(node.parent_index())
+                if parent is not None or node.parent_index() == 0:
+                    new_node = Node(node, parent)
+                    parents[node.index()] = new_node
+                    tree_nodes_count = tree_nodes_count + 1
+                    if node.parent_index() == 0:
+                        root = new_node
+                    break
         return root
 
     def __repr__(self) -> str:
diff --git a/src/parse_tree/parse_tree_item.py b/src/parse_tree/parse_tree_item.py
index 3c0dcd7..79bd4d4 100644
--- a/src/parse_tree/parse_tree_item.py
+++ b/src/parse_tree/parse_tree_item.py
@@ -9,6 +9,9 @@ class ParseItem:
         self._parent_index: int = int(parsed_str[6])
         self._relation: str = parsed_str[7]
 
+    def index(self) -> int:
+        return self._index
+
     def parent_index(self) -> int:
         return self._parent_index