Collect tree to save the context

Author: mmokko

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

@@ -30,6 +30,7 @@
         "y": int,
         "width": int,
         "height": int,
+        "ancestry": int,
     },
 )
 
@@ -53,7 +54,10 @@ def to_window_info(java_window: JavaWindow) -> JavaWindowInfoTypedDict:
 def to_locator_info(context_node: ContextNode) -> LocatorNodeInfoTypedDict:
     ret = {}
     for dct_name in LocatorNodeInfoTypedDict.__annotations__:
-        ret[dct_name] = getattr(context_node.context_info, dct_name)
+        if (dct_name) == "ancestry":
+            ret["ancestry"] = getattr(context_node, dct_name)
+        else:
+            ret[dct_name] = getattr(context_node.context_info, dct_name)
     return cast(LocatorNodeInfoTypedDict, ret)
 
 
@@ -82,7 +86,7 @@ def list_windows(self) -> List[JavaWindowInfoTypedDict]:
         return [to_window_info(window) for window in windows]
 
     def collect_tree(
-        self, window: str, search_depth: int = 8, locator: Optional[str] = None
+        self, window: str, search_depth=1, locator: Optional[str] = None
     ) -> MatchesAndHierarchyTypedDict:
         log.info(f"Collect tree from locator: {locator}")
 

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

@@ -0,0 +1,2 @@
+class NoMatchingLocatorException(Exception):
+    """Match for locator not found."""

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

@@ -13,6 +13,9 @@
 
 
 class ElementInspector:
+    def __init__(self) -> None:
+        self._context: Optional[ContextNode] = None
+
     def _start_event_pump(func, *args, **kwargs):
         def wrapper(self: "ElementInspector", *args, **kwargs):
             from ._event_pump import EventPumpThread
@@ -32,6 +35,47 @@ def wrapper(self: "ElementInspector", *args, **kwargs):
     def list_windows(self, jab_wrapper: JavaAccessBridgeWrapper) -> List[JavaWindow]:
         return jab_wrapper.get_windows()
 
+    def _collect_from_root(
+        self,
+        jab_wrapper: JavaAccessBridgeWrapper,
+        locator: Optional[str],
+        search_depth=1,
+    ) -> ColletedTreeTypedDict:
+        self._context = ContextTree(jab_wrapper, search_depth).root
+        matches: Union[ContextNode, List[ContextNode]] = []
+        if locator:
+            from ._locators import find_elements_from_tree
+
+            matches = find_elements_from_tree(self._context, locator)
+        return {"matches": matches, "tree": self._context}
+
+    def _collect_from_context(
+        self, jab_wrapper: JavaAccessBridgeWrapper, locator: str, search_depth=1
+    ) -> ColletedTreeTypedDict:
+        from threading import RLock
+
+        from ._errors import NoMatchingLocatorException
+        from ._locators import find_elements_from_tree
+
+        # The JavaAccessBridgeWrapper object needs to be inserted into the context as the
+        # object has to be recreated every time we do a new query
+        # TODO: update the ContextTree to introduce the API for this
+        self._context._jab_wrapper = jab_wrapper
+        match = find_elements_from_tree(self._context, locator)
+        node = match[0] if isinstance(match, List) and len(match) > 0 else match
+        if not isinstance(node, ContextNode):
+            raise NoMatchingLocatorException(f"No matching locator for {locator}")
+        self._context = ContextNode(
+            jab_wrapper,
+            node.context,
+            RLock(),
+            node.ancestry,
+            True,
+            search_depth + node.ancestry,
+        )
+        matches = find_elements_from_tree(self._context, locator)
+        return {"matches": matches, "tree": self._context}
+
     @_start_event_pump
     def collect_tree(
         self,
@@ -41,10 +85,8 @@ def collect_tree(
         locator: Optional[str] = None,
     ) -> 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
 
-            matches = find_elements_from_tree(context_tree, locator)
-        return {"matches": matches, "tree": context_tree.root}
+        if not self._context:
+            return self._collect_from_root(jab_wrapper, locator, search_depth)
+        else:
+            return self._collect_from_context(jab_wrapper, locator, search_depth)

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

@@ -22,8 +22,8 @@ def _parse_locator(locator: str, strict_default=False):
         conditions = lvl.split(" and ")
         lvl_search = []
         strict_mode = strict_default
-        for cond in conditions:
-            parts = cond.split(":", 1)
+        for condition in conditions:
+            parts = condition.split(":", 1)
             if len(parts) == 1:
                 parts = ["name", parts[0]]
             elif parts[0].lower() == "strict":