fix: resolve Python imports within scan root

(cherry picked from commit 962a0569b1)

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
 

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

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")