Get tree with matches and hierarchy

Author: mmokko

robocorp-code/src/robocorp_code/inspector/java/java_inspector.py

@@ -4,12 +4,13 @@
 from JABWrapper.jab_wrapper import JavaWindow
 from robocorp_ls_core.robotframework_log import get_logger
 
+from robocorp_code.inspector.java.robocorp_java._inspector import ColletedTreeTypedDict
+
 log = get_logger(__name__)
 
 JavaWindowInfoTypedDict = TypedDict(
     "JavaWindowInfoTypedDict",
     {
-        # Same as JavaWindow
         "pid": int,
         "hwnd": int,
         "title": str,
@@ -32,10 +33,14 @@
     },
 )
 
-LocatorTreeInfoTypedDict = TypedDict(
-    "LocatorTreeInfoTypedDict",
-    {"nodes": List[LocatorNodeInfoTypedDict]},
-)
+
+class MatchesAndHierarchyTypedDict(TypedDict):
+    # A list with the nodes matched (these are the ones that the
+    # locator matched)
+    matched_paths: List[str]
+    # This includes all the entries found along with the full hierarchy
+    # to reach the matched entries.
+    hierarchy: List[LocatorNodeInfoTypedDict]
 
 
 def to_window_info(java_window: JavaWindow) -> JavaWindowInfoTypedDict:
@@ -52,6 +57,18 @@ def to_locator_info(context_node: ContextNode) -> LocatorNodeInfoTypedDict:
     return cast(LocatorNodeInfoTypedDict, ret)
 
 
+def to_matches_and_hierarchy(
+    matches_and_hierarchy: ColletedTreeTypedDict,
+) -> MatchesAndHierarchyTypedDict:
+    matches = (
+        [str(matches_and_hierarchy["matches"])]
+        if type(matches_and_hierarchy["matches"]) == ContextNode
+        else [str(match) for match in matches_and_hierarchy["matches"]]
+    )
+    hierarchy = [to_locator_info(node) for node in matches_and_hierarchy["tree"]]
+    return {"matches": matches, "hierarchy": hierarchy}
+
+
 class JavaInspector:
     def __init__(self):
         from robocorp_code.inspector.java.robocorp_java._inspector import (
@@ -66,8 +83,10 @@ def list_windows(self) -> List[JavaWindowInfoTypedDict]:
 
     def collect_tree(
         self, window: str, search_depth: int = 8, locator: Optional[str] = None
-    ) -> LocatorTreeInfoTypedDict:
+    ) -> MatchesAndHierarchyTypedDict:
         log.info(f"Collect tree from locator: {locator}")
 
-        tree = self._inspector.collect_tree(window, search_depth, locator)
-        return [to_locator_info(node) for node in tree]
+        matches_and_hierarchy = self._inspector.collect_tree(
+            window, search_depth, locator
+        )
+        return to_matches_and_hierarchy(matches_and_hierarchy)

robocorp-code/src/robocorp_code/inspector/java/robocorp_java/_inspector.py

@@ -1,8 +1,16 @@
-from typing import List, Optional, Union
+from typing import List, Optional, TypedDict, Union
 
 from JABWrapper.context_tree import ContextNode, ContextTree
 from JABWrapper.jab_wrapper import JavaAccessBridgeWrapper, JavaWindow
 
+ColletedTreeTypedDict = TypedDict(
+    "ColletedTreeTypedDict",
+    {
+        "matches": Union[ContextNode, List[ContextNode]],
+        "tree": ContextNode,
+    },
+)
+
 
 class ElementInspector:
     def _start_event_pump(func, *args, **kwargs):
@@ -31,11 +39,12 @@ def collect_tree(
         window: str,
         search_depth: int,
         locator: Optional[str] = None,
-    ) -> Union[ContextNode, List[ContextNode]]:
+    ) -> ColletedTreeTypedDict:
         jab_wrapper.switch_window_by_title(window)
         context_tree = ContextTree(jab_wrapper, search_depth)
+        matches: Union[ContextNode, List[ContextNode]] = []
         if locator:
             from ._locators import find_elements_from_tree
 
-            return find_elements_from_tree(context_tree, locator)
-        return context_tree.root
+            matches = find_elements_from_tree(context_tree, locator)
+        return {"matches": matches, "tree": context_tree.root}