Separate uast parents calculation

Signed-off-by: Konstantin Slavnov <[email protected]>

Author: zurk

lookout/style/format/annotations/annotations.py

@@ -1,5 +1,7 @@
 """Annotations for style-analyzer."""
-from typing import Dict, FrozenSet, Optional, Tuple
+from typing import FrozenSet, Optional, Tuple
+
+from lookout.style.format.utils import get_uast_parents
 
 
 class Annotation:
@@ -58,12 +60,26 @@ def from_node(node: "bblfsh.Node") -> "UASTNodeAnnotation":
         return UASTNodeAnnotation(node.start_position.offset, node.end_position.offset, node)
 
 
-# Should be removed when overlapping annotations of one type are allowed.
 class UASTAnnotation(UASTNodeAnnotation):
     """Full UAST of the file annotation."""
 
+    def __init__(self, start: int, stop: int, uast: "bblfsh.Node"):
+        """
+        Create new UASTAnnotation instance.
+
+        Parents mapping for the tree is built during init.
+
+        :param start: Annotation start.
+        :param stop: Annotation end.
+        :param uast: UAST.
+        """
+        super().__init__(start, stop, uast)
+        self._parents = get_uast_parents(uast)
+
     uast = property(lambda self: self._node)
 
+    parents = property(lambda self: self._parents)
+
 
 class TokenAnnotation(Annotation):
     """Annotation for virtual token in the сode."""
@@ -154,21 +170,6 @@ def to_target_annotation(self) -> TargetAnnotation:
         return TargetAnnotation(self.start, self.stop, self._target)
 
 
-class UASTParentsAnnotation(Annotation):
-    """
-    Annotation to store uast parents dict.
-
-    Should be removed as soon as UAST provide parents itself.
-    """
-
-    def __init__(self, start: int, stop: int, parents: Dict[int, "bblfsh.Node"]):
-        """Init."""
-        super().__init__(start, stop)
-        self._parents = parents
-
-    parents = property(lambda self: self._parents)
-
-
 class AccumulatedIndentationAnnotation(Annotation):
     """Annotates accumulated indentation Token."""
 

lookout/style/format/feature_extractor.py

@@ -15,7 +15,7 @@
 
 from lookout.style.format.annotations.annotated_data import AnnotatedData
 from lookout.style.format.annotations.annotations import PathAnnotation, RawTokenAnnotation, \
-    UASTAnnotation, UASTParentsAnnotation
+    UASTAnnotation
 from lookout.style.format.classes import (
     CLASS_INDEX, CLASS_PRINTABLES, CLASS_REPRESENTATIONS, CLS_DOUBLE_QUOTE, CLS_NOOP,
     CLS_SINGLE_QUOTE, CLS_SPACE, CLS_SPACE_DEC, CLS_SPACE_INC, CLS_TAB, CLS_TAB_DEC, CLS_TAB_INC,
@@ -725,19 +725,15 @@ def _parse_file(self, file: AnnotatedData) -> Tuple[List[VirtualNode], Dict[int,
 
         # walk the tree: collect nodes with assigned tokens and build the parents map
         node_tokens = []
-        parents = {}
         queue = [file.get(UASTAnnotation).uast]
         while queue:
             node = queue.pop()
             if node.internal_type in self.node_fixtures:
-                node = self.node_fixtures[node.internal_type](node)
-            for child in node.children:
-                parents[id(child)] = node
+                self.node_fixtures[node.internal_type](node)
             queue.extend(node.children)
             if (node.token or node.start_position and node.end_position
                     and node.start_position != node.end_position and not node.children):
                 node_tokens.append(node)
-        file.add(UASTParentsAnnotation(0, len(file), parents))
         node_tokens.sort(key=lambda n: n.start_position.offset)
         sentinel = bblfsh.Node()
         sentinel.start_position.offset = len(contents)
@@ -845,4 +841,4 @@ def _to_position(raw_lines_data, _lines_start_offset, offset):
             y=None,
         )
         vnodes.append(vnode)
-    return vnodes, file.get(UASTParentsAnnotation).parents
+    return vnodes, file.get(UASTAnnotation).parents

lookout/style/format/langs/javascript/uast_fixers.py

@@ -2,14 +2,13 @@
 import bblfsh
 
 
-def fix_regexp_node(node: bblfsh.Node) -> bblfsh.Node:
+def fix_regexp_node(node: bblfsh.Node) -> None:
     """
     Workaround https://github.com/bblfsh/javascript-driver/issues/37.
 
     Should be removed as soon as issue closed and new driver is used.
     """
     node.token = node.properties["pattern"]
-    return node
 
 
 NODE_FIXTURES = {

lookout/style/format/utils.py

@@ -81,3 +81,20 @@ def get_classification_report(y_pred: numpy.ndarray, y_true: numpy.ndarray,
             "report_full": report_full,
             "confusion_matrix": confusion_matrix,
             "target_names": target_names}
+
+
+def get_uast_parents(uast: "bblfsh.Node") -> Dict[int, "bblfsh.Node"]:
+    """
+    Create a mapping from id of the node in the UAST to its parent Node.
+
+    :param uast: UAST to get parents mapping for.
+    :return: Parents mapping.
+    """
+    parents = {}
+    queue = [uast]
+    while queue:
+        node = queue.pop()
+        for child in node.children:
+            parents[id(child)] = node
+        queue.extend(node.children)
+    return parents