Lint child dependencies recursively (#2226)

## Changes
The current code lints local files without considering parent/child
relationships (where a given 'parent' file imports or runs a 'child'
file).
This makes it impossible to lint files in context, where a child file
uses variables inherited from a parent file.
This PR is progress towards that, by linting child dependencies
recursively.

### Linked issues
Progresses #2155 
Progresses #2156

### Functionality
- [x] modified existing command: `databricks labs ucx lint-local-code`

### Tests
- [x] checked integration tests

---------

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

Author: ericvergnaud

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

@@ -110,6 +110,10 @@ def check_registered_dependency(graph):
     def root(self):
         return self if self._parent is None else self._parent.root
 
+    @property
+    def parent(self):
+        return self._parent
+
     @property
     def all_dependencies(self) -> set[Dependency]:
         dependencies: set[Dependency] = set()
@@ -123,6 +127,28 @@ def add_to_dependencies(graph: DependencyGraph) -> bool:
         self.visit(add_to_dependencies, set())
         return dependencies
 
+    @property
+    def root_dependencies(self) -> set[Dependency]:
+        roots: set[Dependency] = set()
+        children: set[Dependency] = set()
+
+        def add_to_dependencies(graph: DependencyGraph) -> bool:
+            dependency = graph.dependency
+            if dependency in children:
+                return False  # Nothing to do
+            if dependency in roots:
+                roots.remove(dependency)
+                children.add(dependency)
+                return False
+            if graph.parent is None:
+                roots.add(dependency)
+                return False
+            children.add(dependency)
+            return False
+
+        self.visit(add_to_dependencies, set())
+        return roots
+
     @property
     def local_dependencies(self) -> set[Dependency]:
         return {child.dependency for child in self._dependencies.values()}
@@ -136,17 +162,27 @@ def all_paths(self) -> set[Path]:
 
     def all_relative_names(self) -> set[str]:
         """This method is intended to simplify testing"""
-        all_names = set[str]()
-        dependencies = self.all_dependencies
-        for library_root in self._path_lookup.library_roots:
-            for dependency in dependencies:
+        return self._relative_names(self.all_dependencies)
+
+    def _relative_names(self, dependencies: set[Dependency]) -> set[str]:
+        """This method is intended to simplify testing"""
+        names = set[str]()
+        for dependency in dependencies:
+            for library_root in self._path_lookup.library_roots:
                 if not dependency.path.is_relative_to(library_root):
                     continue
                 relative_path = dependency.path.relative_to(library_root).as_posix()
                 if relative_path == ".":
                     continue
-                all_names.add(relative_path)
-        return all_names
+                names.add(relative_path)
+        return names
+
+    def root_paths(self) -> set[Path]:
+        return {d.path for d in self.root_dependencies}
+
+    def root_relative_names(self) -> set[str]:
+        """This method is intended to simplify testing"""
+        return self._relative_names(self.root_dependencies)
 
     # when visit_node returns True it interrupts the visit
     def visit(self, visit_node: Callable[[DependencyGraph], bool | None], visited: set[Path]) -> bool:

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

@@ -105,6 +105,7 @@ def __init__(
     def lint(
         self,
         prompts: Prompts,
+        linted_paths: set[Path],
         path: Path | None,
         stdout: TextIO = sys.stdout,
     ) -> list[LocatedAdvice]:
@@ -116,13 +117,13 @@ def lint(
                 validate=lambda p_: Path(p_).exists(),
             )
             path = Path(response)
-        located_advices = list(self.lint_path(path))
+        located_advices = list(self.lint_path(path, linted_paths))
         for located in located_advices:
             message = located.message_relative_to(path)
             stdout.write(f"{message}\n")
         return located_advices
 
-    def lint_path(self, path: Path) -> Iterable[LocatedAdvice]:
+    def lint_path(self, path: Path, linted_paths: set[Path]) -> Iterable[LocatedAdvice]:
         loader = self._folder_loader if path.is_dir() else self._file_loader
         dependency = Dependency(loader, path)
         graph = DependencyGraph(dependency, None, self._dependency_resolver, self._path_lookup, self._session_state)
@@ -132,15 +133,26 @@ def lint_path(self, path: Path) -> Iterable[LocatedAdvice]:
         for problem in problems:
             problem_path = Path('UNKNOWN') if problem.is_path_missing() else problem.source_path.absolute()
             yield problem.as_advisory().for_path(problem_path)
-        for child_path in graph.all_paths:
-            yield from self._lint_one(child_path)
-
-    def _lint_one(self, path: Path) -> Iterable[LocatedAdvice]:
-        if path.is_dir():
-            return []
-        ctx = self._new_linter_context()
-        linter = FileLinter(ctx, self._path_lookup, self._session_state, path)
-        return [advice.for_path(path) for advice in linter.lint()]
+        for dependency in graph.root_dependencies:
+            yield from self._lint_one(dependency, graph, linted_paths)
+
+    def _lint_one(
+        self, dependency: Dependency, graph: DependencyGraph, linted_paths: set[Path]
+    ) -> Iterable[LocatedAdvice]:
+        if dependency.path in linted_paths:
+            return
+        linted_paths.add(dependency.path)
+        if dependency.path.is_file():
+            ctx = self._new_linter_context()
+            linter = FileLinter(ctx, self._path_lookup, self._session_state, dependency.path)
+            for advice in linter.lint():
+                yield advice.for_path(dependency.path)
+        maybe_graph = graph.locate_dependency(dependency.path)
+        # problems have already been reported while building the graph
+        if maybe_graph.graph:
+            child_graph = maybe_graph.graph
+            for child_dependency in child_graph.local_dependencies:
+                yield from self._lint_one(child_dependency, child_graph, linted_paths)
 
 
 class LocalFileMigrator:

tests/integration/source_code/test_jobs.py

@@ -198,7 +198,7 @@ def test_lint_local_code(simple_ctx):
         light_ctx.dependency_resolver,
         lambda: linter_context,
     )
-    problems = linter.lint(Prompts(), path_to_scan, StringIO())
+    problems = linter.lint(Prompts(), set(), path_to_scan, StringIO())
     assert len(problems) > 0
 
 

tests/unit/source_code/linters/test_files.py

@@ -122,7 +122,9 @@ def test_linter_walks_directory(mock_path_lookup, migration_index):
     linter = LocalCodeLinter(
         file_loader, folder_loader, mock_path_lookup, session_state, resolver, lambda: LinterContext(migration_index)
     )
-    advices = linter.lint(prompts, None)
+    paths: set[Path] = set()
+    advices = linter.lint(prompts, paths, None)
+    assert len(paths) > 10
     assert not advices
 
 
@@ -181,6 +183,7 @@ def test_known_issues(path: Path, migration_index):
         resolver,
         lambda: LinterContext(migration_index, session_state),
     )
-    advices = linter.lint(MockPrompts({}), path)
+    linted: set[Path] = set()
+    advices = linter.lint(MockPrompts({}), linted, path)
     for advice in advices:
         print(repr(advice))

tests/unit/source_code/test_dependencies.py

@@ -34,6 +34,12 @@ def test_dependency_resolver_visits_workspace_notebook_dependencies(simple_depen
     assert maybe.graph.all_relative_names() == {"root3.run.py", "root1.run.py", "leaf1.py", "leaf2.py"}
 
 
+def test_dependency_resolver_locates_root_dependencies(simple_dependency_resolver):
+    maybe = simple_dependency_resolver.build_notebook_dependency_graph(Path("root3.run.py"), CurrentSessionState())
+    assert not maybe.failed
+    assert maybe.graph.root_relative_names() == {"root3.run.py"}
+
+
 def test_dependency_resolver_visits_local_notebook_dependencies(simple_dependency_resolver):
     maybe = simple_dependency_resolver.build_notebook_dependency_graph(Path("root4.py"), CurrentSessionState())
     assert not maybe.failed
@@ -52,7 +58,7 @@ def test_dependency_resolver_visits_workspace_file_dependencies(simple_dependenc
     assert maybe.graph.all_relative_names() == {'leaf1.py', 'leaf2.py', 'root8.py'}
 
 
-def test_dependency_resolver_raises_problem_with_unfound_workspace_notebook_dependency(simple_dependency_resolver):
+def test_dependency_resolver_raises_problem_with_unresolved_workspace_notebook_dependency(simple_dependency_resolver):
     maybe = simple_dependency_resolver.build_notebook_dependency_graph(
         Path("root1-no-leaf.run.py"), CurrentSessionState()
     )
@@ -69,7 +75,7 @@ def test_dependency_resolver_raises_problem_with_unfound_workspace_notebook_depe
     ]
 
 
-def test_dependency_resolver_raises_problem_with_unfound_local_notebook_dependency(simple_dependency_resolver):
+def test_dependency_resolver_raises_problem_with_unresolved_local_notebook_dependency(simple_dependency_resolver):
     maybe = simple_dependency_resolver.build_notebook_dependency_graph(Path("root4-no-leaf.py"), CurrentSessionState())
     assert list(maybe.problems) == [
         DependencyProblem(
@@ -142,14 +148,14 @@ def test_dependency_resolver_terminates_at_known_libraries(empty_index, mock_not
     assert maybe.failed
 
 
-def test_dependency_resolver_raises_problem_with_unfound_root_file(simple_dependency_resolver):
+def test_dependency_resolver_raises_problem_with_unresolved_root_file(simple_dependency_resolver):
     maybe = simple_dependency_resolver.build_local_file_dependency_graph(Path("non-existing.py"), CurrentSessionState())
     assert list(maybe.problems) == [
         DependencyProblem('file-not-found', 'File not found: non-existing.py', Path("non-existing.py"))
     ]
 
 
-def test_dependency_resolver_raises_problem_with_unfound_root_notebook(simple_dependency_resolver):
+def test_dependency_resolver_raises_problem_with_unresolved_root_notebook(simple_dependency_resolver):
     maybe = simple_dependency_resolver.build_notebook_dependency_graph(Path("unknown_notebook"), CurrentSessionState())
     assert list(maybe.problems) == [
         DependencyProblem('notebook-not-found', 'Notebook not found: unknown_notebook', Path("unknown_notebook"))
@@ -183,8 +189,10 @@ def load_dependency(self, path_lookup: PathLookup, dependency: Dependency) -> So
 
     notebook_loader = FailingNotebookLoader()
     notebook_resolver = NotebookResolver(notebook_loader)
-    pip_resolver = PythonLibraryResolver(KnownList())
-    resolver = DependencyResolver(pip_resolver, notebook_resolver, [], mock_path_lookup)
+    known_list = KnownList()
+    pip_resolver = PythonLibraryResolver(known_list)
+    import_resolver = ImportFileResolver(FileLoader(), known_list)
+    resolver = DependencyResolver(pip_resolver, notebook_resolver, import_resolver, mock_path_lookup)
     maybe = resolver.build_notebook_dependency_graph(Path("root5.py"), CurrentSessionState())
     assert list(maybe.problems) == [
         DependencyProblem('cannot-load-notebook', 'Could not load notebook root5.py', Path('<MISSING_SOURCE_PATH>'))