Refactor classes related to Python AST for clarity and consistency (#1854)

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

Author: ericvergnaud

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

@@ -9,13 +9,12 @@
 
 from databricks.labs.ucx.source_code.base import Advisory
 from databricks.labs.ucx.source_code.linters.imports import (
-    ASTLinter,
     DbutilsLinter,
     ImportSource,
-    NodeBase,
     NotebookRunCall,
     SysPathChange,
 )
+from databricks.labs.ucx.source_code.linters.python_ast import Tree, NodeBase
 from databricks.labs.ucx.source_code.path_lookup import PathLookup
 
 
@@ -164,10 +163,10 @@ def visit(self, visit_node: Callable[[DependencyGraph], bool | None], visited: s
 
     def build_graph_from_python_source(self, python_code: str) -> list[DependencyProblem]:
         problems: list[DependencyProblem] = []
-        linter = ASTLinter.parse(python_code)
-        syspath_changes = DbutilsLinter.list_sys_path_changes(linter)
-        run_calls = DbutilsLinter.list_dbutils_notebook_run_calls(linter)
-        import_sources, import_problems = DbutilsLinter.list_import_sources(linter, DependencyProblem)
+        tree = Tree.parse(python_code)
+        syspath_changes = SysPathChange.extract_from_tree(tree)
+        run_calls = DbutilsLinter.list_dbutils_notebook_run_calls(tree)
+        import_sources, import_problems = ImportSource.extract_from_tree(tree, DependencyProblem)
         problems.extend(cast(list[DependencyProblem], import_problems))
         nodes = syspath_changes + run_calls + import_sources
         # need to execute things in intertwined sequence so concat and sort

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

@@ -5,10 +5,10 @@
 from sqlglot.expressions import Table
 
 from databricks.labs.ucx.source_code.base import Advice, Linter, Advisory, Deprecation
-from databricks.labs.ucx.source_code.linters.imports import Visitor, ASTLinter
+from databricks.labs.ucx.source_code.linters.python_ast import Tree, TreeVisitor
 
 
-class DetectDbfsVisitor(Visitor):
+class DetectDbfsVisitor(TreeVisitor):
     """
     Visitor that detects file system paths in Python code and checks them
     against a list of known deprecated paths.
@@ -77,9 +77,9 @@ def lint(self, code: str) -> Iterable[Advice]:
         """
         Lints the code looking for file system paths that are deprecated
         """
-        linter = ASTLinter.parse(code)
+        tree = Tree.parse(code)
         visitor = DetectDbfsVisitor()
-        visitor.visit(linter.root)
+        visitor.visit(tree.root)
         yield from visitor.get_advices()