Infer values from child notebook in magic line (#2091)

## Changes
When linting notebooks, use values from child notebooks loaded via %run
magic line to improve value inference

### Linked issues
Resolves #1201 
Progresses #1901 

### Functionality
None

### Tests
- [x] added unit tests

---------

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

Author: ericvergnaud

pyproject.toml

@@ -97,11 +97,11 @@ path = ".venv"
 test        = "pytest -n 4 --cov src --cov-report=xml --timeout 30 tests/unit --durations 20"
 coverage    = "pytest -n auto --cov src tests/unit --timeout 30 --cov-report=html --durations 20"
 integration = "pytest -n 10 --cov src tests/integration --durations 20"
-fmt         = ["black . --extend-exclude 'tests/unit/source_code/samples/*' --extend-exclude dist",
+fmt         = ["black . --extend-exclude 'tests/unit/source_code/samples/'",
                "ruff check . --fix",
                "mypy --disable-error-code 'annotation-unchecked' --exclude 'tests/unit/source_code/samples/*' --exclude dist .",
                "pylint --output-format=colorized -j 0 src tests"]
-verify      = ["black --check . --extend-exclude 'tests/unit/source_code/samples/*' --extend-exclude dist",
+verify      = ["black --check . --extend-exclude 'tests/unit/source_code/samples/'",
                "ruff check .",
                "mypy --exclude 'tests/unit/source_code/samples/*' --exclude dist .",
                "pylint --output-format=colorized -j 0 src tests"]

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

@@ -6,7 +6,7 @@
 from dataclasses import dataclass
 from pathlib import Path
 
-from astroid import AstroidSyntaxError, NodeNG  # type: ignore
+from astroid import AstroidSyntaxError, Module, NodeNG  # type: ignore
 
 from databricks.sdk.service import compute
 
@@ -205,23 +205,42 @@ def __init__(self, linters: list[PythonLinter]):
 
     def lint(self, code: str) -> Iterable[Advice]:
         try:
-            tree = Tree.normalize_and_parse(code)
-            if self._tree is None:
-                self._tree = tree
-            else:
-                tree = self._tree.append_statements(tree)
-            for linter in self._linters:
-                yield from linter.lint_tree(tree)
+            tree = self._parse_and_append(code)
+            yield from self.lint_tree(tree)
         except AstroidSyntaxError as e:
             yield Failure('syntax-error', str(e), 0, 0, 0, 0)
 
+    def lint_tree(self, tree: Tree) -> Iterable[Advice]:
+        for linter in self._linters:
+            yield from linter.lint_tree(tree)
+
+    def _parse_and_append(self, code: str) -> Tree:
+        tree = Tree.normalize_and_parse(code)
+        self.append_tree(tree)
+        return tree
+
+    def append_tree(self, tree: Tree):
+        if self._tree is None:
+            self._tree = Tree(Module("root"))
+        self._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)
+
+    def append_globals(self, globs: dict):
+        if self._tree is None:
+            self._tree = Tree(Module("root"))
+        self._tree.append_globals(globs)
+
     def process_child_cell(self, code: str):
         try:
             tree = Tree.normalize_and_parse(code)
             if self._tree is None:
                 self._tree = tree
             else:
-                self._tree.append_statements(tree)
+                self._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)

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

@@ -381,16 +381,16 @@ def _lint_task(self, task: jobs.Task, job: jobs.Job):
             if not container:
                 continue
             if isinstance(container, Notebook):
-                yield from self._lint_notebook(container, ctx)
+                yield from self._lint_notebook(container, ctx, session_state)
             if isinstance(container, LocalFile):
-                yield from self._lint_file(container, ctx)
+                yield from self._lint_file(container, ctx, session_state)
 
-    def _lint_file(self, file: LocalFile, ctx: LinterContext):
-        linter = FileLinter(ctx, self._path_lookup, file.path)
+    def _lint_file(self, file: LocalFile, ctx: LinterContext, session_state: CurrentSessionState):
+        linter = FileLinter(ctx, self._path_lookup, session_state, file.path)
         for advice in linter.lint():
             yield file.path, advice
 
-    def _lint_notebook(self, notebook: Notebook, ctx: LinterContext):
-        linter = NotebookLinter(ctx, self._path_lookup, notebook)
+    def _lint_notebook(self, notebook: Notebook, ctx: LinterContext, session_state: CurrentSessionState):
+        linter = NotebookLinter(ctx, self._path_lookup, session_state, notebook)
         for advice in linter.lint():
             yield notebook.path, advice

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

@@ -14,6 +14,7 @@
 from databricks.labs.blueprint.entrypoint import get_logger
 
 from databricks.labs.ucx.hive_metastore.migration_status import MigrationIndex
+from databricks.labs.ucx.source_code.base import CurrentSessionState
 from databricks.labs.ucx.source_code.graph import DependencyProblem
 from databricks.labs.ucx.source_code.linters.context import LinterContext
 from databricks.labs.ucx.source_code.notebooks.sources import FileLinter
@@ -172,8 +173,9 @@ def _analyze_file(cls, known_distributions, library_root, dist_info, module_path
             if module_ref.endswith(suffix):
                 module_ref = module_ref[: -len(suffix)]
         logger.info(f"Processing module: {module_ref}")
-        ctx = LinterContext(empty_index)
-        linter = FileLinter(ctx, PathLookup.from_sys_path(module_path.parent), module_path)
+        session_state = CurrentSessionState()
+        ctx = LinterContext(empty_index, session_state)
+        linter = FileLinter(ctx, PathLookup.from_sys_path(module_path.parent), session_state, module_path)
         known_problems = set()
         for problem in linter.lint():
             known_problems.add(KnownProblem(problem.code, problem.message))

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

@@ -139,7 +139,7 @@ def _lint_one(self, path: Path) -> Iterable[LocatedAdvice]:
         if path.is_dir():
             return []
         ctx = self._new_linter_context()
-        linter = FileLinter(ctx, self._path_lookup, path)
+        linter = FileLinter(ctx, self._path_lookup, self._session_state, path)
         return [advice.for_path(path) for advice in linter.lint()]
 
 

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

@@ -3,6 +3,7 @@
 import abc
 import logging
 from collections.abc import Callable, Iterable
+from pathlib import Path
 from typing import TypeVar, cast
 
 from astroid import (  # type: ignore
@@ -19,6 +20,7 @@
 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.path_lookup import PathLookup
 
 logger = logging.getLogger(__name__)
 
@@ -161,17 +163,18 @@ def __init__(self, node: NodeNG, path: str, is_append: bool):
         self._path = path
         self._is_append = is_append
 
-    @property
-    def node(self):
-        return self._node
-
     @property
     def path(self):
         return self._path
 
-    @property
-    def is_append(self):
-        return self._is_append
+    def apply_to(self, path_lookup: PathLookup):
+        path = Path(self._path)
+        if not path.is_absolute():
+            path = path_lookup.cwd / path
+        if self._is_append:
+            path_lookup.append_path(path)
+            return
+        path_lookup.prepend_path(path)
 
 
 class AbsolutePath(SysPathChange):

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

@@ -210,24 +210,33 @@ def _get_attribute_value(cls, node: Attribute):
             logger.debug(f"Missing handler for {name}")
         return None
 
-    def append_statements(self, tree: Tree) -> Tree:
+    def append_tree(self, tree: Tree) -> Tree:
         if not isinstance(tree.node, Module):
-            raise NotImplementedError(f"Can't append statements from {type(tree.node).__name__}")
+            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 statements to {type(self.node).__name__}")
+            raise NotImplementedError(f"Can't append globals to {type(self.node).__name__}")
         self_module: Module = cast(Module, self.node)
-        for stmt in tree_module.body:
-            stmt.parent = self_module
-            self_module.body.append(stmt)
-        for name, value in tree_module.globals.items():
+        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)
-        # the following may seem strange but it's actually ok to use the original module as tree root
-        return tree
+
+    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

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

@@ -8,6 +8,7 @@
 from collections.abc import Callable
 from enum import Enum
 from pathlib import Path
+from typing import TypeVar
 
 from astroid import Call, Const, ImportFrom, Name, NodeNG  # type: ignore
 from astroid.exceptions import AstroidSyntaxError  # type: ignore
@@ -229,7 +230,8 @@ def is_runnable(self) -> bool:
         return True  # TODO
 
     def build_dependency_graph(self, graph: DependencyGraph) -> list[DependencyProblem]:
-        return PipMagic(self.original_code).build_dependency_graph(graph)
+        node = MagicNode(0, 1, None, end_lineno=0, end_col_offset=len(self.original_code))
+        return PipCommand(node, self.original_code).build_dependency_graph(graph)
 
 
 class CellLanguage(Enum):
@@ -414,10 +416,10 @@ def build_graph_from_python_source(self, python_code: str) -> list[DependencyPro
         import_problems: list[DependencyProblem]
         import_sources, import_problems = ImportSource.extract_from_tree(tree, DependencyProblem.from_node)
         problems.extend(import_problems)
-        magic_commands, command_problems = MagicCommand.extract_from_tree(tree, DependencyProblem.from_node)
+        magic_commands, command_problems = MagicLine.extract_from_tree(tree, DependencyProblem.from_node)
         problems.extend(command_problems)
         nodes = syspath_changes + run_calls + import_sources + magic_commands
-        # need to execute things in intertwined sequence so concat and sort
+        # need to execute things in intertwined sequence so concat and sort them
         for base_node in sorted(nodes, key=lambda node: (node.node.lineno, node.node.col_offset)):
             for problem in self._process_node(base_node):
                 # Astroid line numbers are 1-based.
@@ -437,7 +439,7 @@ def _process_node(self, base_node: NodeBase):
             yield from self._register_notebook(base_node)
         elif isinstance(base_node, ImportSource):
             yield from self._register_import(base_node)
-        elif isinstance(base_node, MagicCommand):
+        elif isinstance(base_node, MagicLine):
             yield from base_node.build_dependency_graph(self._context.parent)
         else:
             logger.warning(f"Can't process {NodeBase.__name__} of type {type(base_node).__name__}")
@@ -466,23 +468,20 @@ def _mutate_path_lookup(self, change: SysPathChange):
                 f"Can't update sys.path from {change.node.as_string()} because the expression cannot be computed",
             )
             return
-        path = Path(change.path)
-        if not path.is_absolute():
-            path = self._context.path_lookup.cwd / path
-        if change.is_append:
-            self._context.path_lookup.append_path(path)
-            return
-        self._context.path_lookup.prepend_path(path)
+        change.apply_to(self._context.path_lookup)
+
 
+T = TypeVar("T")
 
-class MagicCommand(NodeBase):
+
+class MagicLine(NodeBase):
 
     @classmethod
     def extract_from_tree(
-        cls, tree: Tree, problem_factory: Callable[[str, str, NodeNG], DependencyProblem]
-    ) -> tuple[list[MagicCommand], list[DependencyProblem]]:
-        problems: list[DependencyProblem] = []
-        commands: list[MagicCommand] = []
+        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):
@@ -498,36 +497,82 @@ 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 MagicCommand(node, arg.value)
+                yield MagicLine(node, arg.value)
 
     def __init__(self, node: NodeNG, command: bytes):
         super().__init__(node)
         self._command = command.decode()
 
-    def build_dependency_graph(self, graph: DependencyGraph) -> list[DependencyProblem]:
+    def as_magic(self) -> MagicCommand | None:
         if self._command.startswith("%pip") or self._command.startswith("!pip"):
-            cmd = PipMagic(self._command)
-            return cmd.build_dependency_graph(graph)
+            return PipCommand(self.node, self._command)
+        if self._command.startswith("%run"):
+            return RunCommand(self.node, self._command)
+        return None
+
+    def build_dependency_graph(self, graph: DependencyGraph) -> list[DependencyProblem]:
+        magic = self.as_magic()
+        if magic is not None:
+            return magic.build_dependency_graph(graph)
         problem = DependencyProblem.from_node(
             code='unsupported-magic-line', message=f"magic line '{self._command}' is not supported yet", node=self.node
         )
         return [problem]
 
 
-class PipMagic:
+class MagicNode(NodeNG):
+    pass
 
-    def __init__(self, code: str):
+
+class MagicCommand(ABC):
+
+    def __init__(self, node: NodeNG, code: str):
+        self._node = node
         self._code = code
 
+    @abstractmethod
+    def build_dependency_graph(self, graph: DependencyGraph) -> list[DependencyProblem]: ...
+
+
+class RunCommand(MagicCommand):
+
+    def build_dependency_graph(self, graph: DependencyGraph) -> list[DependencyProblem]:
+        path = self.notebook_path
+        if path is not None:
+            problems = graph.register_notebook(path)
+            return [problem.from_node(problem.code, problem.message, self._node) for problem in problems]
+        problem = DependencyProblem.from_node('invalid-run-cell', "Missing notebook path in %run command", self._node)
+        return [problem]
+
+    @property
+    def notebook_path(self) -> Path | None:
+        start = self._code.find(' ')
+        if start < 0:
+            return None
+        path = self._code[start + 1 :].strip().strip('"').strip("'")
+        return Path(path)
+
+
+class PipCommand(MagicCommand):
+
     def build_dependency_graph(self, graph: DependencyGraph) -> list[DependencyProblem]:
         argv = self._split(self._code)
         if len(argv) == 1:
-            return [DependencyProblem("library-install-failed", "Missing command after 'pip'")]
+            return [DependencyProblem.from_node("library-install-failed", "Missing command after 'pip'", self._node)]
         if argv[1] != "install":
-            return [DependencyProblem("library-install-failed", f"Unsupported 'pip' command: {argv[1]}")]
+            return [
+                DependencyProblem.from_node(
+                    "library-install-failed", f"Unsupported 'pip' command: {argv[1]}", self._node
+                )
+            ]
         if len(argv) == 2:
-            return [DependencyProblem("library-install-failed", "Missing arguments after 'pip install'")]
-        return graph.register_library(*argv[2:])  # Skipping %pip install
+            return [
+                DependencyProblem.from_node(
+                    "library-install-failed", "Missing arguments after 'pip install'", self._node
+                )
+            ]
+        problems = graph.register_library(*argv[2:])  # Skipping %pip install
+        return [problem.from_node(problem.code, problem.message, self._node) for problem in problems]
 
     # Cache re-used regex (and ensure issues are raised during class init instead of upon first use).
     _splitter = re.compile(r"(?<!\\)\n")