feat: classify conditional imports

Author: TristanKruse

README.md

@@ -272,6 +272,13 @@ ArchUnitPython detects string-based dynamic imports such as `importlib.import_mo
 from my_app.adapters.sql import Repository  # archunit: ignore
 ```
 
+### Conditional Imports
+
+Imports inside `try` blocks that handle `ImportError` or
+`ModuleNotFoundError` are marked as conditional dependencies. This helps graph
+reports distinguish optional imports and fallback implementations from regular
+runtime imports.
+
 ### Naming Conventions
 
 ```python

src/archunitpython/common/extraction/extract_graph.py

@@ -231,30 +231,48 @@ def _extract_located_imports(file_path: str) -> list[_LocatedImport]:
     imports: list[_LocatedImport] = []
     ignore_directives = _find_ignore_directives(source)
     type_checking_ranges = _find_type_checking_ranges(tree)
+    conditional_import_ranges = _find_conditional_import_ranges(tree)
 
     for node in ast.walk(tree):
         if isinstance(node, ast.Import):
-            is_type = _in_type_checking(node, type_checking_ranges)
-            kind = ImportKind.TYPE_IMPORT if is_type else ImportKind.IMPORT
+            kind = _classify_import(
+                node,
+                ImportKind.IMPORT,
+                type_checking_ranges,
+                conditional_import_ranges,
+            )
             for alias in node.names:
                 imports.append(_LocatedImport(alias.name, kind, node.lineno))
 
         elif isinstance(node, ast.ImportFrom):
-            is_type = _in_type_checking(node, type_checking_ranges)
             if node.level and node.level > 0:
                 # Relative import
-                kind = ImportKind.TYPE_IMPORT if is_type else ImportKind.RELATIVE_IMPORT
+                kind = _classify_import(
+                    node,
+                    ImportKind.RELATIVE_IMPORT,
+                    type_checking_ranges,
+                    conditional_import_ranges,
+                )
                 module = node.module or ""
                 dots = "." * node.level
                 imports.append(_LocatedImport(f"{dots}{module}", kind, node.lineno))
             else:
-                kind = ImportKind.TYPE_IMPORT if is_type else ImportKind.FROM_IMPORT
+                kind = _classify_import(
+                    node,
+                    ImportKind.FROM_IMPORT,
+                    type_checking_ranges,
+                    conditional_import_ranges,
+                )
                 if node.module:
                     imports.append(_LocatedImport(node.module, kind, node.lineno))
 
         elif isinstance(node, ast.Call):
-            is_type = _in_type_checking(node, type_checking_ranges)
-            kind = ImportKind.TYPE_IMPORT if is_type else ImportKind.DYNAMIC_IMPORT
+            kind = _classify_import(
+                node,
+                ImportKind.DYNAMIC_IMPORT,
+                type_checking_ranges,
+                conditional_import_ranges,
+            )
             for module_name in _extract_dynamic_import_names(node):
                 imports.append(_LocatedImport(module_name, kind, node.lineno))
 
@@ -321,6 +339,20 @@ def _extract_dynamic_import_names(node: ast.Call) -> list[str]:
     return []
 
 
+def _classify_import(
+    node: ast.AST,
+    default_kind: ImportKind,
+    type_checking_ranges: list[tuple[int, int]],
+    conditional_import_ranges: list[tuple[int, int]],
+) -> ImportKind:
+    """Classify an import node by special context before syntax kind."""
+    if _in_type_checking(node, type_checking_ranges):
+        return ImportKind.TYPE_IMPORT
+    if _in_conditional_import(node, conditional_import_ranges):
+        return ImportKind.CONDITIONAL_IMPORT
+    return default_kind
+
+
 def _find_type_checking_ranges(tree: ast.Module) -> list[tuple[int, int]]:
     """Find line ranges of TYPE_CHECKING blocks."""
     ranges: list[tuple[int, int]] = []
@@ -346,6 +378,46 @@ def _find_type_checking_ranges(tree: ast.Module) -> list[tuple[int, int]]:
     return sorted(ranges, key=lambda ele: ele[0])
 
 
+def _find_conditional_import_ranges(tree: ast.Module) -> list[tuple[int, int]]:
+    """Find try/except ImportError ranges that contain optional imports."""
+    ranges: list[tuple[int, int]] = []
+
+    for node in ast.walk(tree):
+        if not isinstance(node, ast.Try):
+            continue
+        if not any(_handles_import_error(handler.type) for handler in node.handlers):
+            continue
+
+        ranges.extend(_statement_ranges(node.body))
+        for handler in node.handlers:
+            if _handles_import_error(handler.type):
+                ranges.extend(_statement_ranges(handler.body))
+
+    return sorted(ranges, key=lambda ele: ele[0])
+
+
+def _handles_import_error(node: ast.expr | None) -> bool:
+    """Return True if an except handler catches import-related errors."""
+    if node is None:
+        return False
+    if isinstance(node, ast.Name):
+        return node.id in {"ImportError", "ModuleNotFoundError"}
+    if isinstance(node, ast.Attribute):
+        return node.attr in {"ImportError", "ModuleNotFoundError"}
+    if isinstance(node, ast.Tuple):
+        return any(_handles_import_error(elt) for elt in node.elts)
+    return False
+
+
+def _statement_ranges(statements: list[ast.stmt]) -> list[tuple[int, int]]:
+    """Return line ranges covered by statement blocks."""
+    if not statements:
+        return []
+    start = statements[0].lineno
+    end = max(getattr(statement, "end_lineno", statement.lineno) for statement in statements)
+    return [(start, end)]
+
+
 def _in_type_checking(node: ast.AST, ranges: list[tuple[int, int]]) -> bool:
     """Check if a node is inside a TYPE_CHECKING block."""
     if not hasattr(node, "lineno"):
@@ -354,6 +426,14 @@ def _in_type_checking(node: ast.AST, ranges: list[tuple[int, int]]) -> bool:
     return any(start <= lineno <= end for start, end in ranges)
 
 
+def _in_conditional_import(node: ast.AST, ranges: list[tuple[int, int]]) -> bool:
+    """Check if a node is inside a try/except ImportError block."""
+    if not hasattr(node, "lineno"):
+        return False
+    lineno = node.lineno
+    return any(start <= lineno <= end for start, end in ranges)
+
+
 def _resolve_import(
     import_name: str,
     source_file: str,
@@ -365,7 +445,15 @@ def _resolve_import(
     Returns (resolved_path, is_external).
     The path is normalized with forward slashes.
     """
-    if kind in (ImportKind.RELATIVE_IMPORT, ImportKind.TYPE_IMPORT) and import_name.startswith("."):
+    if (
+        kind
+        in (
+            ImportKind.RELATIVE_IMPORT,
+            ImportKind.TYPE_IMPORT,
+            ImportKind.CONDITIONAL_IMPORT,
+        )
+        and import_name.startswith(".")
+    ):
         # Relative import
         return _resolve_relative_import(import_name, source_file, project_root)
 

src/archunitpython/common/extraction/graph.py

@@ -14,6 +14,7 @@ class ImportKind(Enum):
     RELATIVE_IMPORT = "relative"  # from . import bar / from ..foo import bar
     DYNAMIC_IMPORT = "dynamic"  # __import__('foo') / importlib.import_module()
     TYPE_IMPORT = "type"  # inside TYPE_CHECKING block
+    CONDITIONAL_IMPORT = "conditional"  # inside try/except ImportError
 
 
 @dataclass(frozen=True)

tests/common/test_extract_graph.py

@@ -105,6 +105,32 @@ def test_importlib_import_module(self):
         finally:
             shutil.rmtree(project_root, ignore_errors=True)
 
+    def test_conditional_import(self):
+        temp_root = Path(__file__).resolve().parent / ".tmp"
+        temp_root.mkdir(exist_ok=True)
+        project_root = temp_root / f"project_{uuid4().hex}"
+        project_root.mkdir()
+        file_path = project_root / "loader.py"
+        file_path.write_text(
+            "\n".join(
+                [
+                    "try:",
+                    "    import orjson",
+                    "except ImportError:",
+                    "    import json",
+                    "",
+                ]
+            ),
+            encoding="utf-8",
+        )
+
+        try:
+            imports = _extract_imports(str(file_path))
+            assert ("orjson", ImportKind.CONDITIONAL_IMPORT) in imports
+            assert ("json", ImportKind.CONDITIONAL_IMPORT) in imports
+        finally:
+            shutil.rmtree(project_root, ignore_errors=True)
+
 
 class TestExtractGraph:
     def setup_method(self):
@@ -310,6 +336,75 @@ def test_dynamic_import_resolves_to_internal_edge(self):
         assert ImportKind.DYNAMIC_IMPORT in edges[0].import_kinds
 
 
+class TestConditionalImportGraphHandling:
+    def setup_method(self):
+        clear_graph_cache()
+
+    def _build_conditional_project(self) -> str:
+        temp_root = Path(__file__).resolve().parent / ".tmp"
+        temp_root.mkdir(exist_ok=True)
+        project_root = temp_root / f"project_{uuid4().hex}"
+        project_root.mkdir()
+
+        package_dir = project_root / "sample_project"
+        package_dir.mkdir(parents=True, exist_ok=True)
+
+        (package_dir / "__init__.py").write_text("", encoding="utf-8")
+        (package_dir / "fast_model.py").write_text(
+            "class FastUser:\n    pass\n",
+            encoding="utf-8",
+        )
+        (package_dir / "fallback_model.py").write_text(
+            "class FallbackUser:\n    pass\n",
+            encoding="utf-8",
+        )
+        (package_dir / "service.py").write_text(
+            "\n".join(
+                [
+                    "try:",
+                    "    from sample_project.fast_model import FastUser",
+                    "except ImportError:",
+                    "    from sample_project.fallback_model import FallbackUser",
+                    "",
+                ]
+            ),
+            encoding="utf-8",
+        )
+        self._temp_dir = project_root
+        return str(project_root)
+
+    def teardown_method(self):
+        temp_dir = getattr(self, "_temp_dir", None)
+        if temp_dir is not None:
+            shutil.rmtree(temp_dir, ignore_errors=True)
+
+    def test_import_error_fallback_imports_are_marked_conditional(self):
+        project_root = self._build_conditional_project()
+
+        graph = extract_graph(project_root)
+        service_path = os.path.abspath(
+            os.path.join(project_root, "sample_project", "service.py")
+        ).replace("\\", "/")
+        target_paths = {
+            os.path.abspath(
+                os.path.join(project_root, "sample_project", "fast_model.py")
+            ).replace("\\", "/"),
+            os.path.abspath(
+                os.path.join(project_root, "sample_project", "fallback_model.py")
+            ).replace("\\", "/"),
+        }
+
+        edges = [
+            edge
+            for edge in graph
+            if edge.source == service_path and edge.target in target_paths
+        ]
+
+        assert len(edges) == 2
+        assert all(edge.external is False for edge in edges)
+        assert all(ImportKind.CONDITIONAL_IMPORT in edge.import_kinds for edge in edges)
+
+
 class TestIgnoreDirectives:
     def setup_method(self):
         clear_graph_cache()
@@ -401,3 +496,4 @@ def test_all_kinds_exist(self):
         assert ImportKind.RELATIVE_IMPORT.value == "relative"
         assert ImportKind.DYNAMIC_IMPORT.value == "dynamic"
         assert ImportKind.TYPE_IMPORT.value == "type"
+        assert ImportKind.CONDITIONAL_IMPORT.value == "conditional"