From 1c12006360105f81e9b60ce8f5fdd6d1ed6dbd4b Mon Sep 17 00:00:00 2001 From: Ha1baraA11 <243105435+Ha1baraA11@users.noreply.github.com> Date: Wed, 23 Sep 2026 13:41:53 +0800 Subject: [PATCH] fix: resolve Python imports within scan root (cherry picked from commit 962a0569b12d3f1ed99f63a9bbd7748fb8dd7041) --- graphify/extract.py | 125 +++++++++++++++---------- graphify/extractors/engine.py | 12 ++- tests/test_python_import_resolution.py | 55 +++++++++++ 3 files changed, 139 insertions(+), 53 deletions(-) diff --git a/graphify/extract.py b/graphify/extract.py index b11ebfd5b..56364e020 100644 --- a/graphify/extract.py +++ b/graphify/extract.py @@ -144,6 +144,7 @@ from graphify.extractors.resolution import ( # noqa: E402,F401 _resolve_lua_import_target, _probe_python_module_candidate, _resolve_python_module_path, + _resolve_python_namespace_dir, _resolve_tsconfig_alias, _resolve_workspace_import, _source_key, @@ -181,8 +182,12 @@ def _raise_recursion_limit() -> None: sys.setrecursionlimit(_RECURSION_LIMIT) -def _safe_extract(extractor: Callable, path: Path) -> dict: +def _safe_extract( + extractor: Callable, path: Path, *, scan_root: Path | None = None +) -> dict: try: + if extractor is extract_python: + return extractor(path, root=scan_root) return extractor(path) except RecursionError: print(f" warning: skipped {path} (recursion limit exceeded)", file=sys.stderr, flush=True) @@ -436,15 +441,51 @@ def _resolve_name(node, source: bytes, config: LanguageConfig) -> str | None: # ── Import handlers ─────────────────────────────────────────────────────────── -def _import_python(node, source: bytes, file_nid: str, stem: str, edges: list, str_path: str, scope_stack: list[str] | None = None) -> None: +def _import_python( + node, source: bytes, file_nid: str, stem: str, edges: list, str_path: str, + scope_stack: list[str] | None = None, scan_root: Path | None = None, +) -> None: t = node.type + current_path = Path(str_path) + try: + current_path = current_path.resolve() + except OSError: + pass + root = Path(scan_root) if scan_root is not None else current_path.parent + try: + root = root.resolve() + except OSError: + pass if t == "import_statement": for child in node.children: if child.type in ("dotted_name", "aliased_import"): raw = _read_text(child, source) raw_module, _, raw_alias = raw.partition(" as ") module_name = raw_module.strip().lstrip(".") - tgt_nid = _make_id(module_name) + target_path = _resolve_python_module_path( + module_name, current_path, root, level=0 + ) + # The importer-relative resolver can find a module under one + # nested sys.path root even when the scan contains another file + # with the same dotted name under a different root. Keep those + # cases for the root-wide alias pass, which refuses ambiguous + # aliases; use a direct scan-root hit here. + root_target = _resolve_python_module_path( + module_name, root, root, level=0 + ) + if root_target is None or target_path != root_target: + target_path = None + if target_path is not None: + try: + if target_path.resolve() == current_path: + target_path = None + except OSError: + pass + tgt_nid = ( + _make_id(str(target_path)) + if target_path is not None + else _make_id(module_name) + ) edge = { "source": file_nid, "target": tgt_nid, @@ -455,6 +496,8 @@ def _import_python(node, source: bytes, file_nid: str, stem: str, edges: list, s "source_location": f"L{node.start_point[0] + 1}", "weight": 1.0, } + if target_path is not None: + edge["target_file"] = str(target_path) if raw_alias: # `import pkg.mod as alias` binds the local name `alias`, not # `mod`'s own stem, to the module -- stash it so the cross-file @@ -471,55 +514,35 @@ def _import_python(node, source: bytes, file_nid: str, stem: str, edges: list, s # Relative import - resolve to full path so IDs match file node IDs dots = len(raw) - len(raw.lstrip(".")) module_name = raw.lstrip(".") - base = Path(str_path).parent - for _ in range(dots - 1): - base = base.parent - # A relative import can name a subpackage (a directory with an - # __init__.py), not a module file. Probing the candidate on disk - # (mirroring the companion `imports` edge's - # _resolve_python_module_path) resolves `graphs` -> graphs/__init__.py - # instead of a nonexistent graphs.py: without it the target keeps an - # absolute-path-derived slug that the target_file stamp below can't - # heal, so it dangles per-checkout (#2455). - candidate = base / module_name.replace(".", "/") if module_name else base - resolved = _probe_python_module_candidate(candidate) - if resolved is not None: - target_path = resolved - else: + target_path = _resolve_python_module_path( + module_name, current_path, root, level=dots + ) + if target_path is None: + base = current_path.parent + for _ in range(dots - 1): + base = base.parent rel = (module_name.replace(".", "/") + ".py") if module_name else "__init__.py" target_path = base / rel tgt_nid = _make_id(str(target_path)) else: - # Absolute imports need the same package/file probing as the - # relative arm. Keeping the dotted-name slug here makes - # ``from pkg.sub import thing`` point at ``pkg_sub`` even when - # the real node is ``pkg_sub_init``; the symbol-resolution pass - # may add a second, correct edge, but it cannot repair this - # malformed file-level edge (#3723). Search upward from the - # importer and only probe non-package ancestors, mirroring the - # resolver's sys.path-root rule without requiring the scan root - # at per-file extraction time. - module_rel = raw.replace(".", "/") - current_path = Path(str_path) - try: - current_path = current_path.resolve() - except OSError: - pass - for ancestor in (current_path.parent, *current_path.parent.parents): - if (ancestor / "__init__.py").is_file(): - continue - resolved = _probe_python_module_candidate(ancestor / module_rel) - if resolved is not None: - try: - if resolved.resolve() == current_path: - # An external import can share the importing - # file's basename; never turn that coincidence - # into a self-loop. - continue - except OSError: - pass - target_path = resolved - break + # Use the shared scan-root-aware resolver for absolute imports. + # It stops at the corpus boundary and handles package roots the + # same way as symbol resolution. A namespace package has no file + # node of its own; the corpus pass will emit edges to any + # imported submodules that exist on disk. + target_path = _resolve_python_module_path(raw, current_path, root, level=0) + if target_path is None and _resolve_python_namespace_dir( + raw, current_path, root, level=0 + ) is not None: + return + if target_path is not None: + try: + if target_path.resolve() == current_path: + # Do not turn an import of the current file into a + # self-loop. + target_path = None + except OSError: + pass tgt_nid = _make_id(str(target_path)) if target_path is not None else _make_id(raw) edge = { "source": file_nid, @@ -1726,9 +1749,9 @@ def _normalize_ts_import_types(source: bytes, *, tsx: bool = False) -> bytes | N # ── Public API ──────────────────────────────────────────────────────────────── -def extract_python(path: Path) -> dict: +def extract_python(path: Path, *, root: Path | None = None) -> dict: """Extract classes, functions, and imports from a .py file via tree-sitter AST.""" - result = _extract_generic(path, _PYTHON_CONFIG) + result = _extract_generic(path, _PYTHON_CONFIG, scan_root=root) if "error" not in result: _extract_python_rationale(path, result) return result @@ -6629,7 +6652,7 @@ def _safe_extract_with_xaml_root(extractor, path: Path, root: Path) -> dict: previous_root = _XAML_ACTIVE_EXTRACT_ROOT _XAML_ACTIVE_EXTRACT_ROOT = root.resolve() try: - return _safe_extract(extractor, path) + return _safe_extract(extractor, path, scan_root=root) finally: _XAML_ACTIVE_EXTRACT_ROOT = previous_root diff --git a/graphify/extractors/engine.py b/graphify/extractors/engine.py index 22f669dab..c016961ab 100644 --- a/graphify/extractors/engine.py +++ b/graphify/extractors/engine.py @@ -3459,7 +3459,8 @@ def _lua_is_require_call(node, source: bytes) -> bool: def _extract_generic( - path: Path, config: LanguageConfig, *, source_override: bytes | None = None + path: Path, config: LanguageConfig, *, source_override: bytes | None = None, + scan_root: Path | None = None, ) -> dict: """Generic AST extractor driven by LanguageConfig. @@ -3707,7 +3708,14 @@ def _extract_generic( # Import types if t in config.import_types: if config.import_handler: - imported_modules = config.import_handler(node, source, file_nid, stem, edges, str_path, scope_stack) + if config.ts_module == "tree_sitter_python": + imported_modules = config.import_handler( + node, source, file_nid, stem, edges, str_path, scope_stack, scan_root + ) + else: + imported_modules = config.import_handler( + node, source, file_nid, stem, edges, str_path, scope_stack + ) # Module-level import handlers (Swift) name a module, not a file # path, so there is no pre-existing node to anchor the edge to. # They return (id, label) pairs for which we materialize a diff --git a/tests/test_python_import_resolution.py b/tests/test_python_import_resolution.py index 0efd37e05..3a7d578fa 100644 --- a/tests/test_python_import_resolution.py +++ b/tests/test_python_import_resolution.py @@ -122,6 +122,61 @@ def test_absolute_package_import_targets_package_init(tmp_path: Path): assert "pkg_sub" not in import_targets +def test_plain_absolute_import_targets_package_module(tmp_path: Path): + package_init = _write(tmp_path / "pkg/__init__.py", "") + subpackage_init = _write(tmp_path / "pkg/sub/__init__.py", "") + consumer_path = _write(tmp_path / "app.py", "import pkg.sub\n") + + result = extract( + [package_init, subpackage_init, consumer_path], + cache_root=tmp_path, + root=tmp_path, + parallel=False, + ) + + consumer = _node_id(result, "app.py", "app.py") + subpackage = _node_id(result, "__init__.py", "pkg/sub/__init__.py") + assert _has_edge(result, consumer, subpackage, "imports") + + +def test_absolute_import_does_not_resolve_above_scan_root(tmp_path: Path): + scan_root = tmp_path / "scan" + source = _write(scan_root / "app.py", "from outside_pkg import thing\n") + _write(tmp_path / "outside_pkg/__init__.py", "") + _write(tmp_path / "outside_pkg/thing.py", "def run():\n return 1\n") + + result = extract( + [source], cache_root=tmp_path / "cache", root=scan_root, parallel=False + ) + + app = _node_id(result, "app.py", "app.py") + targets = { + edge["target"] + for edge in result["edges"] + if edge["source"] == app and edge["relation"] == "imports_from" + } + assert targets == {"outside_pkg"} + + +def test_absolute_from_import_keeps_namespace_package_submodule_edge(tmp_path: Path): + namespace_package = tmp_path / "namespace_pkg" + namespace_package.mkdir() + submodule = _write( + namespace_package / "subspace/worker.py", "def run():\n return 1\n" + ) + consumer_path = _write( + tmp_path / "app.py", "from namespace_pkg.subspace import worker\n" + ) + + result = extract( + [consumer_path, submodule], cache_root=tmp_path, root=tmp_path, parallel=False + ) + + consumer = _node_id(result, "app.py", "app.py") + worker = _node_id(result, "worker.py", "namespace_pkg/subspace/worker.py") + assert _has_edge(result, consumer, worker, "imports_from") + + def test_python_package_reexport_resolves_import_and_call_to_origin_symbol(tmp_path: Path): origin = _write(tmp_path / "pkg/foo.py", "def Foo():\n return 1\n") barrel = _write(tmp_path / "pkg/__init__.py", "from .foo import Foo as PublicFoo\n")