enhance python abstract syntax tree API (#2241)

## Changes
Enhance python abstract syntax tree API in preparation of linting with
inherited context
Move Tree static methods to TreeHelper class to avoid 'too many public
methods' linting error

### Linked issues
Progresses #2155 
Progresses #2156 
 
### Functionality
None

### Tests
- [x] added unit tests

---------

Co-authored-by: Eric Vergnaud <[email protected]>

Author: ericvergnaud

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

@@ -6,7 +6,7 @@
 from dataclasses import dataclass
 from pathlib import Path
 
-from astroid import AstroidSyntaxError, Module, NodeNG  # type: ignore
+from astroid import AstroidSyntaxError, NodeNG  # type: ignore
 
 from databricks.sdk.service import compute
 
@@ -221,27 +221,24 @@ def _parse_and_append(self, code: str) -> Tree:
         return tree
 
     def append_tree(self, tree: Tree):
-        if self._tree is None:
-            self._tree = Tree(Module("root"))
-        self._tree.append_tree(tree)
+        self._make_tree().append_tree(tree)
 
     def append_nodes(self, nodes: list[NodeNG]):
-        if self._tree is None:
-            self._tree = Tree(Module("root"))
-        self._tree.append_nodes(nodes)
+        self._make_tree().append_nodes(nodes)
 
     def append_globals(self, globs: dict):
-        if self._tree is None:
-            self._tree = Tree(Module("root"))
-        self._tree.append_globals(globs)
+        self._make_tree().append_globals(globs)
 
     def process_child_cell(self, code: str):
         try:
+            this_tree = self._make_tree()
             tree = Tree.normalize_and_parse(code)
-            if self._tree is None:
-                self._tree = tree
-            else:
-                self._tree.append_tree(tree)
+            this_tree.append_tree(tree)
         except AstroidSyntaxError as e:
             # error already reported when linting enclosing notebook
             logger.warning(f"Failed to parse Python cell: {code}", exc_info=e)
+
+    def _make_tree(self) -> Tree:
+        if self._tree is None:
+            self._tree = Tree.new_module()
+        return self._tree

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

@@ -196,15 +196,15 @@ def visit(self, visit_node: Callable[[DependencyGraph], bool | None], visited: s
                 return True
         return False
 
-    def new_graph_builder_context(self):
-        return GraphBuilderContext(parent=self, path_lookup=self._path_lookup, session_state=self._session_state)
+    def new_dependency_graph_context(self):
+        return DependencyGraphContext(parent=self, path_lookup=self._path_lookup, session_state=self._session_state)
 
     def __repr__(self):
         return f"<DependencyGraph {self.path}>"
 
 
 @dataclass
-class GraphBuilderContext:
+class DependencyGraphContext:
     parent: DependencyGraph
     path_lookup: PathLookup
     session_state: CurrentSessionState

src/databricks/labs/ucx/source_code/linters/files.py

@@ -14,7 +14,7 @@
 from databricks.labs.blueprint.tui import Prompts
 
 from databricks.labs.ucx.source_code.linters.context import LinterContext
-from databricks.labs.ucx.source_code.notebooks.cells import CellLanguage, GraphBuilder
+from databricks.labs.ucx.source_code.notebooks.cells import CellLanguage, PythonCodeAnalyzer
 from databricks.labs.ucx.source_code.graph import (
     BaseImportResolver,
     BaseFileResolver,
@@ -40,9 +40,9 @@ def __init__(self, path: Path, source: str, language: Language):
 
     def build_dependency_graph(self, parent: DependencyGraph) -> list[DependencyProblem]:
         if self._language is CellLanguage.PYTHON:
-            context = parent.new_graph_builder_context()
-            builder = GraphBuilder(context)
-            return builder.build_graph_from_python_source(self._original_code)
+            context = parent.new_dependency_graph_context()
+            analyser = PythonCodeAnalyzer(context, self._original_code)
+            return analyser.build_graph()
         # supported language that does not generate dependencies
         if self._language is CellLanguage.SQL:
             return []

src/databricks/labs/ucx/source_code/linters/pyspark.py

@@ -14,7 +14,7 @@
 )
 from databricks.labs.ucx.source_code.linters.python_infer import InferredValue
 from databricks.labs.ucx.source_code.queries import FromTable
-from databricks.labs.ucx.source_code.linters.python_ast import Tree
+from databricks.labs.ucx.source_code.linters.python_ast import Tree, TreeHelper
 
 
 @dataclass
@@ -61,7 +61,7 @@ def _get_table_arg(self, node: Call):
     def _check_call_context(self, node: Call) -> bool:
         assert isinstance(node.func, Attribute)  # Avoid linter warning
         func_name = node.func.attrname
-        qualified_name = Tree.get_full_function_name(node)
+        qualified_name = TreeHelper.get_full_function_name(node)
 
         # Check if the call_context is None as that means all calls are checked
         if self.call_context is None:

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

@@ -84,6 +84,11 @@ def _convert_magic_lines_to_magic_commands(cls, python_code: str):
                 in_multi_line_comment = not in_multi_line_comment
         return "\n".join(lines)
 
+    @classmethod
+    def new_module(cls):
+        node = Module("root")
+        return Tree(node)
+
     def __init__(self, node: NodeNG):
         self._node: NodeNG = node
 
@@ -118,6 +123,132 @@ def first_statement(self):
                 return self._node.body[0]
         return None
 
+    def __repr__(self):
+        truncate_after = 32
+        code = repr(self._node)
+        if len(code) > truncate_after:
+            code = code[0:truncate_after] + "..."
+        return f"<Tree: {code}>"
+
+    def append_tree(self, tree: Tree) -> Tree:
+        """returns the appended tree, not the consolidated one!"""
+        if not isinstance(tree.node, Module):
+            raise NotImplementedError(f"Can't append tree from {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
+
+    def append_globals(self, globs: dict[str, list[NodeNG]]) -> 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 append_nodes(self, nodes: list[NodeNG]) -> None:
+        if not isinstance(self.node, Module):
+            raise NotImplementedError(f"Can't append statements 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 is_from_module(self, module_name: str) -> bool:
+        # if this is the call's root node, check it against the required module
+        if isinstance(self._node, Name):
+            if self._node.name == module_name:
+                return True
+            root = self.root
+            if not isinstance(root, Module):
+                return False
+            for value in root.globals.get(self._node.name, []):
+                if not isinstance(value, AssignName) or not isinstance(value.parent, Assign):
+                    continue
+                if Tree(value.parent.value).is_from_module(module_name):
+                    return True
+            return False
+        # walk up intermediate calls such as spark.range(...)
+        if isinstance(self._node, Call):
+            return isinstance(self._node.func, Attribute) and Tree(self._node.func.expr).is_from_module(module_name)
+        if isinstance(self._node, Attribute):
+            return Tree(self._node.expr).is_from_module(module_name)
+        return False
+
+    def has_global(self, name: str) -> bool:
+        if not isinstance(self.node, Module):
+            return False
+        self_module: Module = cast(Module, self.node)
+        return self_module.globals.get(name, None) is not None
+
+    def nodes_between(self, first_line: int, last_line: int) -> list[NodeNG]:
+        if not isinstance(self.node, Module):
+            raise NotImplementedError(f"Can't extract nodes from {type(self.node).__name__}")
+        self_module: Module = cast(Module, self.node)
+        nodes: list[NodeNG] = []
+        for node in self_module.body:
+            if node.lineno < first_line:
+                continue
+            if node.lineno > last_line:
+                break
+            nodes.append(node)
+        return nodes
+
+    def globals_between(self, first_line: int, last_line: int) -> dict[str, list[NodeNG]]:
+        if not isinstance(self.node, Module):
+            raise NotImplementedError(f"Can't extract globals from {type(self.node).__name__}")
+        self_module: Module = cast(Module, self.node)
+        globs: dict[str, list[NodeNG]] = {}
+        for key, nodes in self_module.globals.items():
+            nodes_in_scope: list[NodeNG] = []
+            for node in nodes:
+                if node.lineno < first_line or node.lineno > last_line:
+                    continue
+                nodes_in_scope.append(node)
+            if len(nodes_in_scope) > 0:
+                globs[key] = nodes_in_scope
+        return globs
+
+    def line_count(self):
+        if not isinstance(self.node, Module):
+            raise NotImplementedError(f"Can't count lines from {type(self.node).__name__}")
+        self_module: Module = cast(Module, self.node)
+        nodes_count = len(self_module.body)
+        if nodes_count == 0:
+            return 0
+        return 1 + self_module.body[nodes_count - 1].lineno - self_module.body[0].lineno
+
+    def renumber(self, start: int) -> Tree:
+        assert start != 0
+        if not isinstance(self.node, Module):
+            raise NotImplementedError(f"Can't renumber {type(self.node).__name__}")
+        root: Module = self.node
+        # for now renumber in place to avoid the complexity of rebuilding the tree with clones
+
+        def renumber_node(node: NodeNG, offset: int) -> None:
+            for child in node.get_children():
+                renumber_node(child, offset + child.lineno - node.lineno)
+            if node.end_lineno:
+                node.end_lineno = node.end_lineno + offset
+                node.lineno = node.lineno + offset
+
+        nodes = root.body if start > 0 else reversed(root.body)
+        for node in nodes:
+            offset = start - node.lineno
+            renumber_node(node, offset)
+            num_lines = 1 + (node.end_lineno - node.lineno if node.end_lineno else 0)
+            start = start + num_lines if start > 0 else start - num_lines
+        return self
+
+
+class TreeHelper(ABC):
+
     @classmethod
     def extract_call_by_name(cls, call: Call, name: str) -> Call | None:
         """Given a call-chain, extract its sub-call by method name (if it has one)"""
@@ -163,13 +294,6 @@ def is_none(cls, node: NodeNG) -> bool:
             return False
         return node.value is None
 
-    def __repr__(self):
-        truncate_after = 32
-        code = repr(self._node)
-        if len(code) > truncate_after:
-            code = code[0:truncate_after] + "..."
-        return f"<Tree: {code}>"
-
     @classmethod
     def get_full_attribute_name(cls, node: Attribute) -> str:
         return cls._get_attribute_value(node)
@@ -210,55 +334,6 @@ def _get_attribute_value(cls, node: Attribute):
             logger.debug(f"Missing handler for {name}")
         return None
 
-    def append_tree(self, tree: Tree) -> Tree:
-        if not isinstance(tree.node, Module):
-            raise NotImplementedError(f"Can't append tree from {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
-        return tree
-
-    def append_globals(self, globs: dict):
-        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, value in globs.items():
-            statements: list[Expr] = self_module.globals.get(name, None)
-            if statements is None:
-                self_module.globals[name] = list(value)  # clone the source list to avoid side-effects
-                continue
-            statements.extend(value)
-
-    def append_nodes(self, nodes: list[NodeNG]):
-        if not isinstance(self.node, Module):
-            raise NotImplementedError(f"Can't append statements 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 is_from_module(self, module_name: str):
-        # if this is the call's root node, check it against the required module
-        if isinstance(self._node, Name):
-            if self._node.name == module_name:
-                return True
-            root = self.root
-            if not isinstance(root, Module):
-                return False
-            for value in root.globals.get(self._node.name, []):
-                if not isinstance(value, AssignName) or not isinstance(value.parent, Assign):
-                    continue
-                if Tree(value.parent.value).is_from_module(module_name):
-                    return True
-            return False
-        # walk up intermediate calls such as spark.range(...)
-        if isinstance(self._node, Call):
-            return isinstance(self._node.func, Attribute) and Tree(self._node.func.expr).is_from_module(module_name)
-        if isinstance(self._node, Attribute):
-            return Tree(self._node.expr).is_from_module(module_name)
-        return False
-
 
 class TreeVisitor:
 

src/databricks/labs/ucx/source_code/linters/spark_connect.py

@@ -9,7 +9,7 @@
     PythonLinter,
     CurrentSessionState,
 )
-from databricks.labs.ucx.source_code.linters.python_ast import Tree
+from databricks.labs.ucx.source_code.linters.python_ast import Tree, TreeHelper
 
 
 @dataclass
@@ -96,7 +96,7 @@ def _lint_sc(self, node: NodeNG) -> Iterator[Advice]:
             return
         if node.func.attrname not in self._SC_METHODS:
             return
-        function_name = Tree.get_full_function_name(node)
+        function_name = TreeHelper.get_full_function_name(node)
         if not function_name or not function_name.endswith(f"sc.{node.func.attrname}"):
             return
         yield self._rdd_failure(node)
@@ -149,7 +149,7 @@ def _match_sc_set_log_level(self, node: NodeNG) -> Iterator[Advice]:
             return
         if node.func.attrname != 'setLogLevel':
             return
-        function_name = Tree.get_full_function_name(node)
+        function_name = TreeHelper.get_full_function_name(node)
         if not function_name or not function_name.endswith('sc.setLogLevel'):
             return
 
@@ -163,7 +163,7 @@ def _match_sc_set_log_level(self, node: NodeNG) -> Iterator[Advice]:
     def _match_jvm_log(self, node: NodeNG) -> Iterator[Advice]:
         if not isinstance(node, Attribute):
             return
-        attribute_name = Tree.get_full_attribute_name(node)
+        attribute_name = TreeHelper.get_full_attribute_name(node)
         if attribute_name and attribute_name.endswith('org.apache.log4j'):
             yield Failure.from_node(
                 code='spark-logging-in-shared-clusters',
@@ -179,7 +179,7 @@ class UDFMatcher(SharedClusterMatcher):
     def lint(self, node: NodeNG) -> Iterator[Advice]:
         if not isinstance(node, Call):
             return
-        function_name = Tree.get_function_name(node)
+        function_name = TreeHelper.get_function_name(node)
 
         if function_name == 'registerJavaFunction':
             yield Failure.from_node(
@@ -214,7 +214,7 @@ class CatalogApiMatcher(SharedClusterMatcher):
     def lint(self, node: NodeNG) -> Iterator[Advice]:
         if not isinstance(node, Attribute):
             return
-        if node.attrname == 'catalog' and Tree.get_full_attribute_name(node).endswith('spark.catalog'):
+        if node.attrname == 'catalog' and TreeHelper.get_full_attribute_name(node).endswith('spark.catalog'):
             yield Failure.from_node(
                 code='catalog-api-in-shared-clusters',
                 message=f'spark.catalog functions require DBR 14.3 LTS or above on {self._cluster_type_str()}',
@@ -226,7 +226,7 @@ class CommandContextMatcher(SharedClusterMatcher):
     def lint(self, node: NodeNG) -> Iterator[Advice]:
         if not isinstance(node, Call):
             return
-        function_name = Tree.get_full_function_name(node)
+        function_name = TreeHelper.get_full_function_name(node)
         if function_name and function_name.endswith('getContext.toJson'):
             yield Failure.from_node(
                 code='to-json-in-shared-clusters',