From 1b8f4b56b1daa8c3d6e8d31aa6b92a5b9aee113c Mon Sep 17 00:00:00 2001
From: Cor Zuurmond <jczuurmond@protonmail.com>
Date: Fri, 17 Jan 2025 10:22:54 +0100
Subject: [PATCH] Pass run cell's tree as parent to the notebook it is running

---
 .../labs/ucx/source_code/notebooks/sources.py | 22 +++++++++++--------
 1 file changed, 13 insertions(+), 9 deletions(-)

diff --git a/src/databricks/labs/ucx/source_code/notebooks/sources.py b/src/databricks/labs/ucx/source_code/notebooks/sources.py
index eb6eea63e1..915e7e729c 100644
--- a/src/databricks/labs/ucx/source_code/notebooks/sources.py
+++ b/src/databricks/labs/ucx/source_code/notebooks/sources.py
@@ -184,7 +184,7 @@ def lint(self) -> Iterable[Advice]:
         return
 
     def _load_tree_from_notebook(
-        self, notebook: Notebook, register_trees: bool, *, parent_tree: Tree | None = None
+        self, notebook: Notebook, register_trees: bool, *, parent_tree: Tree | None
     ) -> Failure | None:
         for cell in notebook.cells:
             failure = None
@@ -198,7 +198,7 @@ def _load_tree_from_notebook(
         return None
 
     def _load_tree_from_python_cell(
-        self, python_cell: PythonCell, register_trees: bool, parent_tree: Tree | None = None
+        self, python_cell: PythonCell, register_trees: bool, *, parent_tree: Tree | None
     ) -> Failure | None:
         maybe_tree = Tree.maybe_normalized_parse(python_cell.original_code)
         if maybe_tree.failure:
@@ -224,7 +224,7 @@ def _load_children_from_tree(self, tree: Tree) -> Failure | None:
         # need to execute things in intertwined sequence so concat and sort them
         nodes = list(cast(Module, tree.node).body)
         base_nodes = sorted(base_nodes, key=lambda node: (node.node.lineno, node.node.col_offset))
-        return self._load_children_with_base_nodes(nodes, base_nodes)
+        return self._load_children_with_base_nodes(nodes, base_nodes, parent_tree=tree)
 
     @staticmethod
     def _list_run_magic_lines(tree: Tree) -> Iterable[MagicLine]:
@@ -237,24 +237,28 @@ def _ignore_problem(_code: str, _message: str, _node: NodeNG) -> None:
             if isinstance(command.as_magic(), RunCommand):
                 yield command
 
-    def _load_children_with_base_nodes(self, nodes: list[NodeNG], base_nodes: list[NodeBase]) -> Failure | None:
+    def _load_children_with_base_nodes(
+        self, nodes: list[NodeNG], base_nodes: list[NodeBase], *, parent_tree: Tree | None
+    ) -> Failure | None:
         for base_node in base_nodes:
-            failure = self._load_children_with_base_node(nodes, base_node)
+            failure = self._load_children_with_base_node(nodes, base_node, parent_tree=parent_tree)
             if failure:
                 return failure
         return None
 
-    def _load_children_with_base_node(self, nodes: list[NodeNG], base_node: NodeBase) -> Failure | None:
+    def _load_children_with_base_node(
+        self, nodes: list[NodeNG], base_node: NodeBase, *, parent_tree: Tree | None
+    ) -> Failure | None:
         while len(nodes) > 0:
             node = nodes.pop(0)
             if node.lineno < base_node.node.lineno:
                 continue
-            failure = self._load_children_from_base_node(base_node)
+            failure = self._load_children_from_base_node(base_node, parent_tree=parent_tree)
             if failure:
                 return failure
         return None
 
-    def _load_children_from_base_node(self, base_node: NodeBase) -> Failure | None:
+    def _load_children_from_base_node(self, base_node: NodeBase, *, parent_tree: Tree | None) -> Failure | None:
         if isinstance(base_node, SysPathChange):
             failure = self._mutate_path_lookup(base_node)
             if failure:
@@ -270,7 +274,7 @@ def _load_children_from_base_node(self, base_node: NodeBase) -> Failure | None:
                     node=base_node.node,
                 )
                 return failure
-            return self._load_tree_from_notebook(notebook, False)
+            return self._load_tree_from_notebook(notebook, False, parent_tree=parent_tree)
         return None
 
     def _mutate_path_lookup(self, change: SysPathChange) -> Failure | None: