From 5f0f600e02d82249e0735738b43cc75317eff790 Mon Sep 17 00:00:00 2001 From: TristanKruse Date: Sun, 5 Jul 2026 00:35:42 +0200 Subject: [PATCH] feat: classify conditional imports --- README.md | 7 ++ .../common/extraction/extract_graph.py | 104 ++++++++++++++++-- src/archunitpython/common/extraction/graph.py | 1 + tests/common/test_extract_graph.py | 96 ++++++++++++++++ 4 files changed, 200 insertions(+), 8 deletions(-) diff --git a/README.md b/README.md index 9cca0bc..fd42f55 100644 --- a/README.md +++ b/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 diff --git a/src/archunitpython/common/extraction/extract_graph.py b/src/archunitpython/common/extraction/extract_graph.py index b28194e..5c822c5 100644 --- a/src/archunitpython/common/extraction/extract_graph.py +++ b/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) diff --git a/src/archunitpython/common/extraction/graph.py b/src/archunitpython/common/extraction/graph.py index 79e7a61..5ace535 100644 --- a/src/archunitpython/common/extraction/graph.py +++ b/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) diff --git a/tests/common/test_extract_graph.py b/tests/common/test_extract_graph.py index e73a072..b976452 100644 --- a/tests/common/test_extract_graph.py +++ b/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"