Refactor python code analysis (#2521)

## Changes
Python code analysis is under linter package which is misleading since
these are general utilities.
This PR reorganizes code in preparation of PR #2519 

### Linked issues
Progresses #2350

### Functionality
None

### Tests
- [x] ran unit tests

---------

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

Author: ericvergnaud

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

@@ -15,7 +15,7 @@
 from databricks.sdk.service.workspace import Language
 
 from databricks.labs.blueprint.paths import WorkspacePath
-from databricks.labs.ucx.source_code.linters.python_ast import Tree
+from databricks.labs.ucx.source_code.python.python_ast import Tree
 
 # Code mapping between LSP, PyLint, and our own diagnostics:
 # | LSP                       | PyLint     | Our            |

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

@@ -11,7 +11,7 @@
     NodeNG,
 )
 from databricks.labs.ucx.source_code.base import Advisory, CurrentSessionState, is_a_notebook
-from databricks.labs.ucx.source_code.linters.python_ast import Tree
+from databricks.labs.ucx.source_code.python.python_ast import Tree
 from databricks.labs.ucx.source_code.path_lookup import PathLookup
 
 logger = logging.Logger(__name__)

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

@@ -31,7 +31,7 @@
     DependencyGraphWalker,
 )
 from databricks.labs.ucx.source_code.linters.context import LinterContext
-from databricks.labs.ucx.source_code.linters.python_ast import Tree
+from databricks.labs.ucx.source_code.python.python_ast import Tree
 from databricks.labs.ucx.source_code.notebooks.sources import FileLinter
 from databricks.labs.ucx.source_code.path_lookup import PathLookup
 

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

@@ -12,8 +12,8 @@
     PythonLinter,
     SqlLinter,
 )
-from databricks.labs.ucx.source_code.linters.python_ast import Tree, TreeVisitor
-from databricks.labs.ucx.source_code.linters.python_infer import InferredValue
+from databricks.labs.ucx.source_code.python.python_ast import Tree, TreeVisitor
+from databricks.labs.ucx.source_code.python.python_infer import InferredValue
 
 logger = logging.getLogger(__name__)
 

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

@@ -7,7 +7,7 @@
 from typing import TextIO
 
 from databricks.labs.ucx.source_code.base import LocatedAdvice, CurrentSessionState, file_language, is_a_notebook
-from databricks.labs.ucx.source_code.linters.python_ast import Tree
+from databricks.labs.ucx.source_code.python.python_ast import Tree
 from databricks.labs.ucx.source_code.notebooks.loaders import NotebookLoader
 from databricks.labs.ucx.source_code.notebooks.sources import FileLinter
 from databricks.labs.ucx.source_code.path_lookup import PathLookup

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

@@ -18,8 +18,8 @@
 )
 
 from databricks.labs.ucx.source_code.base import Advice, Advisory, CurrentSessionState, PythonLinter
-from databricks.labs.ucx.source_code.linters.python_ast import Tree, NodeBase, TreeVisitor
-from databricks.labs.ucx.source_code.linters.python_infer import InferredValue
+from databricks.labs.ucx.source_code.python.python_ast import Tree, NodeBase, TreeVisitor
+from databricks.labs.ucx.source_code.python.python_infer import InferredValue
 from databricks.labs.ucx.source_code.path_lookup import PathLookup
 
 logger = logging.getLogger(__name__)

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

@@ -12,9 +12,9 @@
     CurrentSessionState,
     PythonLinter,
 )
-from databricks.labs.ucx.source_code.linters.python_infer import InferredValue
+from databricks.labs.ucx.source_code.python.python_infer import InferredValue
 from databricks.labs.ucx.source_code.queries import FromTableSqlLinter
-from databricks.labs.ucx.source_code.linters.python_ast import Tree, TreeHelper
+from databricks.labs.ucx.source_code.python.python_ast import Tree, TreeHelper
 
 
 @dataclass

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

@@ -11,7 +11,7 @@
 )
 from databricks.sdk.service.compute import DataSecurityMode
 
-from databricks.labs.ucx.source_code.linters.python_ast import Tree, TreeHelper
+from databricks.labs.ucx.source_code.python.python_ast import Tree, TreeHelper
 
 
 @dataclass

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

@@ -9,7 +9,7 @@
     Advice,
     PythonLinter,
 )
-from databricks.labs.ucx.source_code.linters.python_ast import Tree, TreeHelper
+from databricks.labs.ucx.source_code.python.python_ast import Tree, TreeHelper
 
 
 @dataclass

src/databricks/labs/ucx/source_code/notebooks/cells.py

@@ -5,13 +5,10 @@
 import shlex
 from abc import ABC, abstractmethod
 from ast import parse as parse_python
-from collections.abc import Callable, Iterable
 from enum import Enum
 from pathlib import Path
-from typing import TypeVar, cast
 
-from astroid import Call, Const, ImportFrom, Name, NodeNG, Try  # type: ignore
-from astroid.exceptions import AstroidSyntaxError  # type: ignore
+from astroid import NodeNG  # type: ignore
 from sqlglot import parse as parse_sql, ParseError as SQLParseError
 
 from databricks.sdk.service.workspace import Language
@@ -23,14 +20,8 @@
     DependencyGraphContext,
     InheritedContext,
 )
-from databricks.labs.ucx.source_code.linters.imports import (
-    SysPathChange,
-    DbutilsPyLinter,
-    ImportSource,
-    NotebookRunCall,
-    UnresolvedPath,
-)
-from databricks.labs.ucx.source_code.linters.python_ast import Tree, NodeBase
+from databricks.labs.ucx.source_code.python.python_analyzer import PythonCodeAnalyzer
+from databricks.labs.ucx.source_code.notebooks.magic import MagicNode, MagicCommand, magic_command_factory
 
 # use a specific logger for sqlglot warnings so we can disable them selectively
 sqlglot_logger = logging.getLogger(f"{__name__}.sqlglot")
@@ -403,242 +394,15 @@ def wrap_with_magic(self, code: str, cell_language: CellLanguage) -> str:
         return "\n".join(lines)
 
 
-class PythonCodeAnalyzer:
-
-    def __init__(self, context: DependencyGraphContext, python_code: str):
-        self._context = context
-        self._python_code = python_code
-
-    def build_graph(self) -> list[DependencyProblem]:
-        """Check python code for dependency-related problems.
-
-        Returns:
-            A list of dependency problems; position information is relative to the python code itself.
-        """
-        problems: list[DependencyProblem] = []
-        try:
-            _, nodes, parse_problems = self._parse_and_extract_nodes()
-            problems.extend(parse_problems)
-        except AstroidSyntaxError as e:
-            logger.debug(f"Could not parse Python code: {self._python_code}", exc_info=True)
-            problems.append(DependencyProblem('parse-error', f"Could not parse Python code: {e}"))
-            return problems
-        for base_node in nodes:
-            for problem in self._build_graph_from_node(base_node):
-                # Astroid line numbers are 1-based.
-                problem = problem.replace(
-                    start_line=base_node.node.lineno - 1,
-                    start_col=base_node.node.col_offset,
-                    end_line=(base_node.node.end_lineno or 1) - 1,
-                    end_col=base_node.node.end_col_offset or 0,
-                )
-                problems.append(problem)
-        return problems
-
-    def build_inherited_context(self, child_path: Path) -> InheritedContext:
-        try:
-            tree, nodes, _ = self._parse_and_extract_nodes()
-        except AstroidSyntaxError:
-            logger.debug(f"Could not parse Python code: {self._python_code}", exc_info=True)
-            return InheritedContext(None, False)
-        if len(nodes) == 0:
-            return InheritedContext(tree, False)
-        context = InheritedContext(Tree.new_module(), False)
-        last_line = -1
-        for base_node in nodes:
-            # append nodes
-            node_line = base_node.node.lineno
-            nodes = tree.nodes_between(last_line + 1, node_line - 1)
-            context.tree.append_nodes(nodes)
-            globs = tree.globals_between(last_line + 1, node_line - 1)
-            context.tree.append_globals(globs)
-            last_line = node_line
-            # process node
-            child_context = self._build_inherited_context_from_node(base_node, child_path)
-            context = context.append(child_context, True)
-            if context.found:
-                return context
-        line_count = tree.line_count()
-        if last_line < line_count:
-            nodes = tree.nodes_between(last_line + 1, line_count)
-            context.tree.append_nodes(nodes)
-            globs = tree.globals_between(last_line + 1, line_count)
-            context.tree.append_globals(globs)
-        return context
-
-    def _parse_and_extract_nodes(self) -> tuple[Tree, list[NodeBase], Iterable[DependencyProblem]]:
-        problems: list[DependencyProblem] = []
-        tree = Tree.normalize_and_parse(self._python_code)
-        syspath_changes = SysPathChange.extract_from_tree(self._context.session_state, tree)
-        run_calls = DbutilsPyLinter.list_dbutils_notebook_run_calls(tree)
-        import_sources: list[ImportSource]
-        import_problems: list[DependencyProblem]
-        import_sources, import_problems = ImportSource.extract_from_tree(tree, DependencyProblem.from_node)
-        problems.extend(import_problems)
-        magic_lines, command_problems = MagicLine.extract_from_tree(tree, DependencyProblem.from_node)
-        problems.extend(command_problems)
-        # need to evaluate things in intertwined sequence so concat and sort them
-        nodes: list[NodeBase] = cast(list[NodeBase], syspath_changes + run_calls + import_sources + magic_lines)
-        nodes = sorted(nodes, key=lambda node: (node.node.lineno, node.node.col_offset))
-        return tree, nodes, problems
-
-    def _build_graph_from_node(self, base_node: NodeBase) -> Iterable[DependencyProblem]:
-        if isinstance(base_node, SysPathChange):
-            yield from self._mutate_path_lookup(base_node)
-        elif isinstance(base_node, NotebookRunCall):
-            yield from self._register_notebook(base_node)
-        elif isinstance(base_node, ImportSource):
-            yield from self._register_import(base_node)
-        elif isinstance(base_node, MagicLine):
-            yield from base_node.build_dependency_graph(self._context.parent)
-        else:
-            logger.warning(f"Can't build graph for node {NodeBase.__name__} of type {type(base_node).__name__}")
-
-    def _build_inherited_context_from_node(self, base_node: NodeBase, child_path: Path) -> InheritedContext:
-        if isinstance(base_node, SysPathChange):
-            self._mutate_path_lookup(base_node)
-            return InheritedContext(None, False)
-        if isinstance(base_node, ImportSource):
-            # nothing to do, Astroid takes care of imports
-            return InheritedContext(None, False)
-        if isinstance(base_node, NotebookRunCall):
-            # nothing to do, dbutils.notebook.run uses a dedicated context
-            return InheritedContext(None, False)
-        if isinstance(base_node, MagicLine):
-            return base_node.build_inherited_context(self._context, child_path)
-        logger.warning(f"Can't build inherited context for node {NodeBase.__name__} of type {type(base_node).__name__}")
-        return InheritedContext(None, False)
-
-    def _register_import(self, base_node: ImportSource) -> Iterable[DependencyProblem]:
-        prefix = ""
-        if isinstance(base_node.node, ImportFrom) and base_node.node.level is not None:
-            prefix = "." * base_node.node.level
-        name = base_node.name or ""
-        problems = self._context.parent.register_import(prefix + name)
-        for problem in problems:
-            prob = self._filter_import_problem_in_try_except(problem, base_node)
-            if prob is not None:
-                yield prob
-
-    @classmethod
-    def _filter_import_problem_in_try_except(
-        cls, problem: DependencyProblem, base_node: ImportSource
-    ) -> DependencyProblem | None:
-        if problem.code != 'import-not-found':
-            return problem
-        # is base_node in a try-except clause ?
-        node = base_node.node.parent
-        while node and not isinstance(node, Try):
-            node = node.parent
-        if cls._is_try_except_import_error(node):
-            return None
-        return problem
-
-    @classmethod
-    def _is_try_except_import_error(cls, node: Try | None) -> bool:
-        if not isinstance(node, Try):
-            return False
-        for handler in node.handlers:
-            if isinstance(handler.type, Name):
-                if handler.type.name == "ImportError":
-                    return True
-        return False
-
-    def _register_notebook(self, base_node: NotebookRunCall) -> Iterable[DependencyProblem]:
-        has_unresolved, paths = base_node.get_notebook_paths(self._context.session_state)
-        if has_unresolved:
-            yield DependencyProblem(
-                'dependency-cannot-compute-value',
-                f"Can't check dependency from {base_node.node.as_string()} because the expression cannot be computed",
-            )
-        for path in paths:
-            # notebooks ran via dbutils.notebook.run do not inherit or propagate context
-            yield from self._context.parent.register_notebook(Path(path), False)
-
-    def _mutate_path_lookup(self, change: SysPathChange) -> Iterable[DependencyProblem]:
-        if isinstance(change, UnresolvedPath):
-            yield DependencyProblem(
-                'sys-path-cannot-compute-value',
-                f"Can't update sys.path from {change.node.as_string()} because the expression cannot be computed",
-            )
-            return
-        change.apply_to(self._context.path_lookup)
-
-
-T = TypeVar("T")
-
-
-class MagicLine(NodeBase):
-
-    @classmethod
-    def extract_from_tree(
-        cls, tree: Tree, problem_factory: Callable[[str, str, NodeNG], T]
-    ) -> tuple[list[MagicLine], list[T]]:
-        problems: list[T] = []
-        commands: list[MagicLine] = []
-        try:
-            nodes = tree.locate(Call, [("magic_command", Name)])
-            for command in cls._make_commands_for_magic_command_call_nodes(nodes):
-                commands.append(command)
-        except Exception as e:  # pylint: disable=broad-except
-            logger.debug(f"Internal error while checking magic commands in tree: {tree.root}", exc_info=True)
-            problem = problem_factory('internal-error', f"While checking magic commands: {e}", tree.root)
-            problems.append(problem)
-        return commands, problems
+class RunCommand(MagicCommand):
 
-    @classmethod
-    def _make_commands_for_magic_command_call_nodes(cls, nodes: list[Call]):
-        for node in nodes:
-            arg = node.args[0]
-            if isinstance(arg, Const):
-                yield MagicLine(node, arg.value)
-
-    def __init__(self, node: NodeNG, command: bytes):
-        super().__init__(node)
-        self._command = command.decode()
-
-    def as_magic(self) -> MagicCommand | None:
-        if self._command.startswith("%pip") or self._command.startswith("!pip"):
-            return PipCommand(self.node, self._command)
-        if self._command.startswith("%run"):
-            return RunCommand(self.node, self._command)
+    @staticmethod
+    @magic_command_factory
+    def factory(command: str, node: NodeNG) -> MagicCommand | None:
+        if command.startswith("%run"):
+            return RunCommand(node, command)
         return None
 
-    def build_dependency_graph(self, parent: DependencyGraph) -> list[DependencyProblem]:
-        magic = self.as_magic()
-        if magic is not None:
-            return magic.build_dependency_graph(parent)
-        problem = DependencyProblem.from_node(
-            code='unsupported-magic-line', message=f"magic line '{self._command}' is not supported yet", node=self.node
-        )
-        return [problem]
-
-    def build_inherited_context(self, context: DependencyGraphContext, child_path: Path) -> InheritedContext:
-        magic = self.as_magic()
-        if magic is not None:
-            return magic.build_inherited_context(context, child_path)
-        return InheritedContext(None, False)
-
-
-class MagicNode(NodeNG):
-    pass
-
-
-class MagicCommand(ABC):
-
-    def __init__(self, node: NodeNG, code: str):
-        self._node = node
-        self._code = code
-
-    @abstractmethod
-    def build_dependency_graph(self, parent: DependencyGraph) -> list[DependencyProblem]: ...
-
-    def build_inherited_context(self, _context: DependencyGraphContext, _child_path: Path) -> InheritedContext:
-        return InheritedContext(None, False)
-
-
-class RunCommand(MagicCommand):
-
     def build_dependency_graph(self, parent: DependencyGraph) -> list[DependencyProblem]:
         path = self.notebook_path
         if path is not None:
@@ -678,6 +442,13 @@ def build_inherited_context(self, context: DependencyGraphContext, child_path: P
 
 class PipCommand(MagicCommand):
 
+    @staticmethod
+    @magic_command_factory
+    def factory(command: str, node: NodeNG) -> MagicCommand | None:
+        if command.startswith("%pip") or command.startswith("!pip"):
+            return PipCommand(node, command)
+        return None
+
     def build_dependency_graph(self, parent: DependencyGraph) -> list[DependencyProblem]:
         argv = self._split(self._code)
         if len(argv) == 1: