Rename Python AST's Tree methods for clarity (#3524)

## Changes
Rename Python AST's `Tree` methods for clarity

- Add docstrings
- Rename `append_` to `attach_` and `extend_` to more precisely describe
what the methods do
- Let `attach_` and `extend_` always return `None`
- Extend unit testing

### Linked issues

Progresses #3514
Precedes #3520

### Functionality

- [x] modified code linting related
- [x] modified existing command: `databricks labs ucx lint-local-code`

### Tests

- [x] manually tested
- [x] added and modified unit tests

Author: JCZuurmond

src/databricks/labs/ucx/source_code/graph.py

@@ -559,7 +559,7 @@ def append(self, context: InheritedContext, copy_found: bool) -> InheritedContex
         if tree is None:
             return InheritedContext(self.tree, found, context.problems)
         new_tree = self.tree or Tree.new_module()
-        new_tree.append_tree(tree)
+        new_tree.attach_child_tree(tree)
         new_problems = itertools.chain(self.problems, context.problems)
         return InheritedContext(new_tree, found, new_problems)
 

src/databricks/labs/ucx/source_code/jobs.py

@@ -662,7 +662,7 @@ def _collect_from_notebook(
                     logger.warning(maybe_tree.failure.message)
                     continue
                 assert maybe_tree.tree is not None
-                inherited_tree.append_tree(maybe_tree.tree)
+                inherited_tree.attach_child_tree(maybe_tree.tree)
 
     def _collect_from_source(
         self,

src/databricks/labs/ucx/source_code/python/python_analyzer.py

@@ -73,9 +73,9 @@ def build_inherited_context(self, child_path: Path) -> InheritedContext:
             # append nodes
             node_line = base_node.node.lineno
             nodes = tree.nodes_between(last_line + 1, node_line - 1)
-            context.tree.append_nodes(nodes)
+            context.tree.attach_nodes(nodes)
             globs = tree.globals_between(last_line + 1, node_line - 1)
-            context.tree.append_globals(globs)
+            context.tree.extend_globals(globs)
             last_line = node_line
             # process node
             child_context = self._build_inherited_context_from_node(base_node, child_path)
@@ -86,9 +86,9 @@ def build_inherited_context(self, child_path: Path) -> InheritedContext:
         assert context.tree is not None, "Tree should be initialized"
         if last_line < line_count:
             nodes = tree.nodes_between(last_line + 1, line_count)
-            context.tree.append_nodes(nodes)
+            context.tree.attach_nodes(nodes)
             globs = tree.globals_between(last_line + 1, line_count)
-            context.tree.append_globals(globs)
+            context.tree.extend_globals(globs)
         return context
 
     def _parse_and_extract_nodes(self) -> tuple[Tree | None, list[NodeBase], Iterable[DependencyProblem]]:

src/databricks/labs/ucx/source_code/python/python_ast.py

@@ -16,7 +16,6 @@
     Call,
     ClassDef,
     Const,
-    Expr,
     Import,
     ImportFrom,
     Instance,
@@ -211,36 +210,45 @@ def __repr__(self):
             code = code[0:truncate_after] + "..."
         return f"<Tree: {code}>"
 
-    def append_tree(self, tree: Tree) -> Tree:
-        """returns the appended tree, not the consolidated one!"""
+    def attach_child_tree(self, tree: Tree) -> None:
+        """Attach a child tree.
+
+        Attaching a child tree is a **stateful** operation for both the parent and child tree. After attaching a child
+        tree, a tree can be traversed starting from the parent or child tree. From both starting points all nodes in
+        both trees can be reached, though, the order of nodes will be different as that is relative to the starting
+        point.
+        """
         if not isinstance(tree.node, Module):
-            raise NotImplementedError(f"Can't append tree from {type(tree.node).__name__}")
+            raise NotImplementedError(f"Cannot attach child tree: {type(tree.node).__name__}")
         tree_module: Module = cast(Module, tree.node)
-        self.append_nodes(tree_module.body)
-        self.append_globals(tree_module.globals)
-        # the following may seem strange but it's actually ok to use the original module as tree root
-        # because each node points to the correct parent (practically, the tree is now only a list of statements)
-        return tree
+        self.attach_nodes(tree_module.body)
+        self.extend_globals(tree_module.globals)
 
-    def append_globals(self, globs: dict[str, list[Expr]]) -> None:
-        if not isinstance(self.node, Module):
-            raise NotImplementedError(f"Can't append globals to {type(self.node).__name__}")
-        self_module: Module = cast(Module, self.node)
-        for name, values in globs.items():
-            statements: list[Expr] = self_module.globals.get(name, None)
-            if statements is None:
-                self_module.globals[name] = list(values)  # clone the source list to avoid side-effects
-                continue
-            statements.extend(values)
+    def attach_nodes(self, nodes: list[NodeNG]) -> None:
+        """Attach nodes.
 
-    def append_nodes(self, nodes: list[NodeNG]) -> None:
+        Attaching nodes is a **stateful** operation for both this tree's node, the parent node, and the child nodes.
+        After attaching the nodes, the parent node has the nodes in its body and the child nodes have this tree's node
+        as parent node.
+        """
         if not isinstance(self.node, Module):
-            raise NotImplementedError(f"Can't append statements to {type(self.node).__name__}")
+            raise NotImplementedError(f"Cannot attach nodes to: {type(self.node).__name__}")
         self_module: Module = cast(Module, self.node)
         for node in nodes:
             node.parent = self_module
             self_module.body.append(node)
 
+    def extend_globals(self, globs: dict[str, list[NodeNG]]) -> None:
+        """Extend globals by extending the global values for each global key.
+
+        Extending globals is a stateful operation for this `Tree` (`self`), similarly to extending a list.
+        """
+        if not isinstance(self.node, Module):
+            raise NotImplementedError(f"Cannot extend globals to: {type(self.node).__name__}")
+        self_module: Module = cast(Module, self.node)
+        for global_key, global_values in globs.items():
+            self_module.globals[global_key] = self_module.globals.get(global_key, []) + global_values
+
     def is_instance_of(self, class_name: str) -> bool:
         for inferred in self.node.inferred():
             if inferred is Uninferable:
@@ -693,13 +701,13 @@ def _parse_and_append(self, code: str) -> MaybeTree:
         return maybe_tree
 
     def append_tree(self, tree: Tree) -> None:
-        self._make_tree().append_tree(tree)
+        self._make_tree().attach_child_tree(tree)
 
     def append_nodes(self, nodes: list[NodeNG]) -> None:
-        self._make_tree().append_nodes(nodes)
+        self._make_tree().attach_nodes(nodes)
 
     def append_globals(self, globs: dict) -> None:
-        self._make_tree().append_globals(globs)
+        self._make_tree().extend_globals(globs)
 
     def process_child_cell(self, code: str) -> None:
         this_tree = self._make_tree()
@@ -709,7 +717,7 @@ def process_child_cell(self, code: str) -> None:
             logger.warning(maybe_tree.failure.message)
             return
         assert maybe_tree.tree is not None
-        this_tree.append_tree(maybe_tree.tree)
+        this_tree.attach_child_tree(maybe_tree.tree)
 
     def collect_dfsas(self, source_code: str) -> Iterable[DirectFsAccess]:
         maybe_tree = self._parse_and_append(source_code)

tests/unit/source_code/python/test_python_ast.py

@@ -1,5 +1,6 @@
 import pytest
-from astroid import Assign, Attribute, Call, Const, Expr, Module, Name  # type: ignore
+import astroid  # type: ignore
+from astroid import Assign, AssignName, Attribute, Call, Const, Expr, Module, Name  # type: ignore
 
 from databricks.labs.ucx.source_code.python.python_ast import Tree, TreeHelper
 from databricks.labs.ucx.source_code.python.python_infer import InferredValue
@@ -139,21 +140,27 @@ def test_ignores_magic_marker_in_multiline_comment() -> None:
     assert True
 
 
-def test_appends_statements() -> None:
-    source_1 = "a = 'John'"
-    maybe_tree_1 = Tree.maybe_normalized_parse(source_1)
-    assert maybe_tree_1.tree is not None, maybe_tree_1.failure
-    tree_1 = maybe_tree_1.tree
-    source_2 = 'b = f"Hello {a}!"'
-    maybe_tree_2 = Tree.maybe_normalized_parse(source_2)
-    assert maybe_tree_2.tree is not None, maybe_tree_2.failure
-    tree_2 = maybe_tree_2.tree
-    tree_3 = tree_1.append_tree(tree_2)
-    nodes = tree_3.locate(Assign, [])
-    tree = Tree(nodes[0].value)  # tree_3 only contains tree_2 statements
-    values = list(InferredValue.infer_from_node(tree.node))
-    strings = list(value.as_string() for value in values)
-    assert strings == ["Hello John!"]
+def test_tree_attach_child_tree_infers_value() -> None:
+    """Attaching trees allows traversing from both parent and child."""
+    inferred_string = "Hello John!"
+    parent_source, child_source = "a = 'John'", 'b = f"Hello {a}!"'
+    parent_maybe_tree = Tree.maybe_normalized_parse(parent_source)
+    child_maybe_tree = Tree.maybe_normalized_parse(child_source)
+
+    assert parent_maybe_tree.tree is not None, parent_maybe_tree.failure
+    assert child_maybe_tree.tree is not None, child_maybe_tree.failure
+
+    parent_maybe_tree.tree.attach_child_tree(child_maybe_tree.tree)
+
+    nodes = parent_maybe_tree.tree.locate(Assign, [])
+    tree = Tree(nodes[1].value)  # Starting from the parent, we are looking for the last assign
+    strings = [value.as_string() for value in InferredValue.infer_from_node(tree.node)]
+    assert strings == [inferred_string]
+
+    nodes = child_maybe_tree.tree.locate(Assign, [])
+    tree = Tree(nodes[0].value)  # Starting from child, we are looking for the first assign
+    strings = [value.as_string() for value in InferredValue.infer_from_node(tree.node)]
+    assert strings == [inferred_string]
 
 
 def test_is_from_module() -> None:
@@ -194,28 +201,23 @@ def test_is_instance_of(source, name, class_name) -> None:
     assert Tree(var[0]).is_instance_of(class_name)
 
 
-def test_supports_recursive_refs_when_checking_module() -> None:
-    source_1 = """
-    df = spark.read.csv("hi")
-    """
-    source_2 = """
-    df = df.withColumn(stuff)
-    """
-    source_3 = """
-    df = df.withColumn(stuff2)
-    """
-    maybe_tree = Tree.maybe_normalized_parse(source_1)
-    assert maybe_tree.tree is not None, maybe_tree.failure
-    main_tree = maybe_tree.tree
-    maybe_tree_2 = Tree.maybe_normalized_parse(source_2)
-    assert maybe_tree_2.tree is not None, maybe_tree_2.failure
-    tree_2 = maybe_tree_2.tree
-    main_tree.append_tree(tree_2)
-    maybe_tree_3 = Tree.maybe_normalized_parse(source_3)
-    assert maybe_tree_3.tree is not None, maybe_tree_3.failure
-    tree_3 = maybe_tree_3.tree
-    main_tree.append_tree(tree_3)
-    assign = tree_3.locate(Assign, [])[0]
+def test_tree_attach_child_tree_propagates_module_reference() -> None:
+    """The spark module should propagate from the parent tree."""
+    source_1 = "df = spark.read.csv('hi')"
+    source_2 = "df = df.withColumn(stuff)"
+    source_3 = "df = df.withColumn(stuff2)"
+    first_line_maybe_tree = Tree.maybe_normalized_parse(source_1)
+    second_line_maybe_tree = Tree.maybe_normalized_parse(source_2)
+    third_line_maybe_tree = Tree.maybe_normalized_parse(source_3)
+
+    assert first_line_maybe_tree.tree, first_line_maybe_tree.failure
+    assert second_line_maybe_tree.tree, second_line_maybe_tree.failure
+    assert third_line_maybe_tree.tree, third_line_maybe_tree.failure
+
+    first_line_maybe_tree.tree.attach_child_tree(second_line_maybe_tree.tree)
+    first_line_maybe_tree.tree.attach_child_tree(third_line_maybe_tree.tree)
+
+    assign = third_line_maybe_tree.tree.locate(Assign, [])[0]
     assert Tree(assign.value).is_from_module("spark")
 
 
@@ -302,6 +304,53 @@ def test_is_builtin(source, name, is_builtin) -> None:
     assert False  # could not locate call
 
 
+def test_tree_attach_nodes_sets_parent() -> None:
+    node = astroid.extract_node("b = a + 2")
+    maybe_tree = Tree.maybe_normalized_parse("a = 1")
+    assert maybe_tree.tree, maybe_tree.failure
+
+    maybe_tree.tree.attach_nodes([node])
+
+    assert node.parent == maybe_tree.tree.node
+
+
+def test_tree_attach_nodes_adds_node_to_body() -> None:
+    node = astroid.extract_node("b = a + 2")
+    maybe_tree = Tree.maybe_normalized_parse("a = 1")
+    assert maybe_tree.tree, maybe_tree.failure
+
+    maybe_tree.tree.attach_nodes([node])
+
+    assert maybe_tree.tree.node.body[-1] == node
+
+
+def test_tree_extend_globals_adds_assign_name_to_tree() -> None:
+    maybe_tree = Tree.maybe_normalized_parse("a = 1")
+    assert maybe_tree.tree, maybe_tree.failure
+
+    node = astroid.extract_node("b = a + 2")
+    assign_name = next(node.get_children())
+    assert isinstance(assign_name, AssignName)
+
+    maybe_tree.tree.extend_globals({"b": [assign_name]})
+
+    assert isinstance(maybe_tree.tree.node, Module)
+    assert maybe_tree.tree.node.globals.get("b") == [assign_name]
+
+
+def test_tree_attach_child_tree_appends_globals_to_parent_tree() -> None:
+    parent_tree = Tree.maybe_normalized_parse("a = 1")
+    child_tree = Tree.maybe_normalized_parse("b = a + 2")
+
+    assert parent_tree.tree, parent_tree.failure
+    assert child_tree.tree, child_tree.failure
+
+    parent_tree.tree.attach_child_tree(child_tree.tree)
+
+    assert set(parent_tree.tree.node.globals.keys()) == {"a", "b"}
+    assert set(child_tree.tree.node.globals.keys()) == {"b"}
+
+
 def test_first_statement_is_none() -> None:
     node = Const("xyz")
     assert not Tree(node).first_statement()
@@ -311,14 +360,19 @@ def test_repr_is_truncated() -> None:
     assert len(repr(Tree(Const("xyz")))) <= (32 + len("...") + len("<Tree: >"))
 
 
-def test_append_tree_fails() -> None:
-    with pytest.raises(NotImplementedError):
-        Tree(Const("xyz")).append_tree(Tree(Const("xyz")))
+def test_tree_attach_child_tree_raises_not_implemented_error_for_constant_node() -> None:
+    with pytest.raises(NotImplementedError, match="Cannot attach child tree: .*"):
+        Tree(Const("xyz")).attach_child_tree(Tree(Const("xyz")))
 
 
-def test_append_node_fails() -> None:
-    with pytest.raises(NotImplementedError):
-        Tree(Const("xyz")).append_nodes([])
+def test_tree_attach_nodes_raises_not_implemented_error_for_constant_node() -> None:
+    with pytest.raises(NotImplementedError, match="Cannot attach nodes to: .*"):
+        Tree(Const("xyz")).attach_nodes([])
+
+
+def test_extend_globals_raises_not_implemented_error_for_constant_node() -> None:
+    with pytest.raises(NotImplementedError, match="Cannot extend globals to: .*"):
+        Tree(Const("xyz")).extend_globals({})
 
 
 def test_nodes_between_fails() -> None:
@@ -330,11 +384,6 @@ def test_has_global_fails() -> None:
     assert not Tree.new_module().has_global("xyz")
 
 
-def test_append_globals_fails() -> None:
-    with pytest.raises(NotImplementedError):
-        Tree(Const("xyz")).append_globals({})
-
-
 def test_globals_between_fails() -> None:
     with pytest.raises(NotImplementedError):
         Tree(Const("xyz")).line_count()