Filter out missing import problems for imports within a try-except clause with ImportError (#2332)

## Changes
Filter out missing import problems for imports within a try-except
clause with ImportError

### Linked issues
Resolves #1705 

### Functionality
None

### Tests
- [x] added unit tests

---------

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

Author: ericvergnaud

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

@@ -10,7 +10,7 @@
 from pathlib import Path
 from typing import TypeVar, cast
 
-from astroid import Call, Const, ImportFrom, Name, NodeNG  # type: ignore
+from astroid import Call, Const, ImportFrom, Name, NodeNG, Try  # type: ignore
 from astroid.exceptions import AstroidSyntaxError  # type: ignore
 from sqlglot import parse as parse_sql, ParseError as SQLParseError
 
@@ -482,7 +482,7 @@ def _parse_and_extract_nodes(self) -> tuple[Tree, list[NodeBase], Iterable[Depen
         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):
+    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):
@@ -509,14 +509,42 @@ def _build_inherited_context_from_node(self, base_node: NodeBase, child_path: Pa
         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):
+    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 ""
-        yield from self._context.parent.register_import(prefix + name)
+        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
 
-    def _register_notebook(self, base_node: NotebookRunCall):
+    @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(
@@ -527,7 +555,7 @@ def _register_notebook(self, base_node: NotebookRunCall):
             # 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):
+    def _mutate_path_lookup(self, change: SysPathChange) -> Iterable[DependencyProblem]:
         if isinstance(change, UnresolvedPath):
             yield DependencyProblem(
                 'sys-path-cannot-compute-value',

tests/unit/source_code/linters/test_python_imports.py

@@ -1,12 +1,16 @@
 from __future__ import annotations
 
+from pathlib import Path
+
 import pytest
 
 from databricks.labs.ucx.source_code.base import CurrentSessionState
-from databricks.labs.ucx.source_code.graph import DependencyProblem
+from databricks.labs.ucx.source_code.graph import Dependency, DependencyGraph, DependencyProblem
+from databricks.labs.ucx.source_code.linters.files import FileLoader
 
 from databricks.labs.ucx.source_code.linters.imports import DbutilsLinter, ImportSource, SysPathChange
 from databricks.labs.ucx.source_code.linters.python_ast import Tree
+from databricks.labs.ucx.source_code.notebooks.cells import PythonCodeAnalyzer
 
 
 def test_linter_returns_empty_list_of_dbutils_notebook_run_calls():
@@ -27,27 +31,27 @@ def test_linter_returns_list_of_dbutils_notebook_run_calls():
 
 def test_linter_returns_empty_list_of_imports():
     tree = Tree.parse('')
-    assert not ImportSource.extract_from_tree(tree, DependencyProblem)[0]
+    assert not ImportSource.extract_from_tree(tree, DependencyProblem.from_node)[0]
 
 
 def test_linter_returns_import():
     tree = Tree.parse('import x')
-    assert ["x"] == [node.name for node in ImportSource.extract_from_tree(tree, DependencyProblem)[0]]
+    assert ["x"] == [node.name for node in ImportSource.extract_from_tree(tree, DependencyProblem.from_node)[0]]
 
 
 def test_linter_returns_import_from():
     tree = Tree.parse('from x import z')
-    assert ["x"] == [node.name for node in ImportSource.extract_from_tree(tree, DependencyProblem)[0]]
+    assert ["x"] == [node.name for node in ImportSource.extract_from_tree(tree, DependencyProblem.from_node)[0]]
 
 
 def test_linter_returns_import_module():
     tree = Tree.parse('importlib.import_module("x")')
-    assert ["x"] == [node.name for node in ImportSource.extract_from_tree(tree, DependencyProblem)[0]]
+    assert ["x"] == [node.name for node in ImportSource.extract_from_tree(tree, DependencyProblem.from_node)[0]]
 
 
 def test_linter_returns__import__():
     tree = Tree.parse('importlib.__import__("x")')
-    assert ["x"] == [node.name for node in ImportSource.extract_from_tree(tree, DependencyProblem)[0]]
+    assert ["x"] == [node.name for node in ImportSource.extract_from_tree(tree, DependencyProblem.from_node)[0]]
 
 
 def test_linter_returns_appended_absolute_paths():
@@ -191,3 +195,17 @@ def test_infers_dbutils_notebook_run_dynamic_value(code, expected) -> None:
         _, paths = call.get_notebook_paths(CurrentSessionState())
         all_paths.extend(paths)
     assert all_paths == expected
+
+
+def test_failing_import_in_try_except_is_silent(simple_dependency_resolver, mock_path_lookup) -> None:
+    code = """
+try:
+    import missing_library
+except ImportError as e:
+    pass
+"""
+    dep = Dependency(FileLoader(), Path(""))
+    graph = DependencyGraph(dep, None, simple_dependency_resolver, mock_path_lookup, CurrentSessionState())
+    analyzer = PythonCodeAnalyzer(graph.new_dependency_graph_context(), code)
+    problems = analyzer.build_graph()
+    assert not problems