From 33edd89b43a746f33a911ad00a8e4bdf25a68ebf Mon Sep 17 00:00:00 2001 From: Alex Tumanov Date: Sat, 3 Oct 2026 14:36:16 -0500 Subject: [PATCH] fix(extract): preserve Python import bindings across sibling cases Use existing plain-import provenance to distinguish separate and same-line from-imports. Reject type aliases and dynamic namespace references, and protect changed and unchanged external callers through incremental CLI extraction. Fixes #3793 AI-Assisted-By: OpenAI Codex (cherry picked from commit 5e240c8bd39f58b87d7840ae40c1e34488781375) --- graphify/extract.py | 16 +++--- graphify/extractors/engine.py | 33 +++++++---- tests/test_cache.py | 1 + tests/test_extract.py | 78 ++++++++++++++++++-------- tests/test_python_import_resolution.py | 25 ++++++++- 5 files changed, 110 insertions(+), 43 deletions(-) diff --git a/graphify/extract.py b/graphify/extract.py index e2a1d0209..b463ae388 100644 --- a/graphify/extract.py +++ b/graphify/extract.py @@ -338,6 +338,8 @@ def _suppress_ambiguous_python_imports( module_name = edge.pop("_python_import_module", None) bindings = edge.pop("_python_import_bindings", []) is_module_binding = edge.pop("_python_import_module_binding", False) + if is_module_binding is True: + edge["_python_plain_import"] = True marker_only = edge.pop("_python_import_marker_only", False) if not module_name: if not marker_only: @@ -4083,10 +4085,6 @@ def _resolve_python_member_calls( # to the module in the importing file. Keyed by (importing file, target module) # so two files aliasing the same module differently each match their own. import_alias_by_filenode: dict[str, dict[str, str]] = {} - from_targets = {(e.get("source"), e.get("target")) for e in all_edges - if e.get("relation") == "imports_from"} - from_import_sites = {(e.get("source"), e.get("source_file"), e.get("source_location")) - for e in all_edges if e.get("relation") == "imports_from"} external_plain_imports: dict[str, dict[str, list[str]]] = {} for e in all_edges: if e.get("relation") in ("imports", "imports_from"): @@ -4094,11 +4092,9 @@ def _resolve_python_member_calls( alias = e.get("local_alias") if alias: import_alias_by_filenode.setdefault(e.get("source"), {})[e.get("target")] = _key(alias) - if e.get("relation") == "imports": + if e.get("relation") == "imports" and e.get("_python_plain_import") is True: src, target = e.get("source"), e.get("target") - if ((src, target) in from_targets - or (src, e.get("source_file"), e.get("source_location")) in from_import_sites - or not isinstance(target, str)): + if not isinstance(src, str) or not isinstance(target, str): continue if target in node_by_id or target in contains_children or target in file_of_node: continue @@ -4150,7 +4146,8 @@ def _resolve_python_member_calls( # older cached facts, from-imports, and ambiguous bindings fail closed. if rc.get("_python_receiver_shadowed") is False: caller_file = file_of_node.get(caller) - targets = external_plain_imports.get(caller_file, {}).get(receiver, []) + targets = (external_plain_imports.get(caller_file, {}).get(receiver, []) + if caller_file is not None else []) if len(targets) == 1: _emit_call(caller, targets[0], rc) continue @@ -9020,6 +9017,7 @@ def extract( # so it cannot be popped at that earlier point without breaking the fix. for e in all_edges: e.pop("local_alias", None) + e.pop("_python_plain_import", None) # Tag AST provenance so the incremental watch rebuild can distinguish # AST-extracted nodes from semantic/LLM nodes. On a full re-extraction diff --git a/graphify/extractors/engine.py b/graphify/extractors/engine.py index ba5422773..0d00dd767 100644 --- a/graphify/extractors/engine.py +++ b/graphify/extractors/engine.py @@ -1547,20 +1547,17 @@ def _python_receiver_shadow_checker(root, source: bytes): # A nested function can mutate module/captured bindings later. Dynamic # namespace operations can also replace an import without an assignment AST. dynamic_binding = False + dynamic_names = {"exec", "eval", "globals", "locals", "vars", "setattr", "delattr"} pending = [root] while pending: candidate = pending.pop() - if candidate.type in {"global_statement", "nonlocal_statement"}: + # Namespace mutators can be called directly, qualified, or saved under + # an alias. Reject references as well as calls, including imported aliases. + if (candidate.type in {"global_statement", "nonlocal_statement"} + or (candidate.type == "identifier" + and _read_text(candidate, source) in dynamic_names)): dynamic_binding = True break - if candidate.type == "call": - function = candidate.child_by_field_name("function") - if (function is not None and function.type == "identifier" - and _read_text(function, source) in { - "exec", "eval", "globals", "locals", "vars", "setattr", "delattr", - }): - dynamic_binding = True - break pending.extend(candidate.named_children) def target_names(node, names: set[str]) -> bool: @@ -1578,7 +1575,13 @@ def _python_receiver_shadow_checker(root, source: bytes): def import_names(node) -> list[str] | None: names: list[str] = [] - for item in node.named_children: + past_import = False + for item in node.children: + if item.type == "import": + past_import = True + continue + if not past_import: + continue if item.type == "aliased_import": alias = item.child_by_field_name("alias") if alias is None: @@ -1596,7 +1599,8 @@ def _python_receiver_shadow_checker(root, source: bytes): return scopes[key] bound: set[str] = set() plain_imports: dict[str, int] = {} - unsafe = bool(scope.has_error or (scope.type == "module" and dynamic_binding)) + unsafe = bool(scope.has_error or scope.child_by_field_name("type_parameters") + or (scope.type == "module" and dynamic_binding)) if scope.type in {"function_definition", "lambda"}: bound.update(_python_param_names(scope.child_by_field_name("parameters"), source)) @@ -1629,6 +1633,13 @@ def _python_receiver_shadow_checker(root, source: bytes): if kind in {"assignment", "augmented_assignment", "for_statement", "for_in_clause"}: if not target_names(child.child_by_field_name("left"), bound): unsafe = True + elif kind == "type_alias_statement": + left = child.child_by_field_name("left") + name = _read_text(left, source).split("[", 1)[0].strip() if left is not None else "" + if name.isidentifier(): + bound.add(name) + else: + unsafe = True elif kind == "named_expression": if not target_names(child.child_by_field_name("name"), bound): unsafe = True diff --git a/tests/test_cache.py b/tests/test_cache.py index 534e1dfd3..f011c434a 100644 --- a/tests/test_cache.py +++ b/tests/test_cache.py @@ -665,6 +665,7 @@ def test_python_receiver_shadow_schema_retires_same_version_raw_calls( assert any(rc.get("_python_receiver_shadowed") is False for rc in fresh.get("raw_calls", [])) warm = extract([target], cache_root=tmp_path, root=tmp_path) + assert len([e for e in cold["edges"] if e["relation"] == "calls"]) == 1 assert [e for e in cold["edges"] if e["relation"] == "calls"] == [ e for e in warm["edges"] if e["relation"] == "calls" ] diff --git a/tests/test_extract.py b/tests/test_extract.py index 182aef9ec..ef3b10b21 100644 --- a/tests/test_extract.py +++ b/tests/test_extract.py @@ -2024,21 +2024,24 @@ def test_python_external_resolver_requires_explicit_false_shadow_fact(hint): nodes = [{"id": "caller", "label": "caller.py", "source_file": "caller.py"}, {"id": "caller_fetch", "label": "fetch()", "source_file": "caller.py"}] edges = [{"source": "caller", "target": "caller_fetch", "relation": "contains"}, - {"source": "caller", "target": "requests", "relation": "imports"}] + {"source": "caller", "target": "requests", "relation": "imports", + "_python_plain_import": True}] _resolve_python_member_calls([{"raw_calls": [raw]}], nodes, edges) assert not any(e["relation"] == "calls" for e in edges) -@pytest.mark.parametrize(("receiver", "imports", "extra_edges"), [ - ("Requests", [("requests", "requests")], []), - ("Request", [("requests", None)], []), - ("requests", [], []), - ("requests", [("requests", None), ("other", "requests")], []), - ("requests", [("requests", None)], [("imports_from", "requests")]), - ("requests", [("requests", None)], [("imports_from", "other")]), +@pytest.mark.parametrize(("receiver", "imports", "producer_hint"), [ + ("Requests", [("requests", "requests")], True), + ("Request", [("requests", None)], True), + ("requests", [], True), + ("requests", [("requests", None), ("other", "requests")], True), + ("requests", [("requests", None)], None), + ("requests", [("requests", None)], False), + ("requests", [("requests", None)], 1), + ("requests", [("requests", None)], "true"), ]) def test_python_external_resolver_requires_unique_exact_plain_import( - receiver, imports, extra_edges, + receiver, imports, producer_hint, ): from graphify.extract import _resolve_python_member_calls @@ -2050,10 +2053,8 @@ def test_python_external_resolver_requires_unique_exact_plain_import( edges = [{"source": "caller", "target": "caller_fetch", "relation": "contains"}] edges.extend({"source": "caller", "target": target, "relation": "imports", "source_file": "caller.py", "source_location": "L1", + **({"_python_plain_import": producer_hint} if producer_hint is not None else {}), **({"local_alias": alias} if alias else {})} for target, alias in imports) - edges.extend({"source": "caller", "target": target, "relation": relation, - "source_file": "caller.py", "source_location": "L1"} - for relation, target in extra_edges) _resolve_python_member_calls([{"raw_calls": [raw]}], nodes, edges) assert not any(e["relation"] == "calls" for e in edges) @@ -2069,7 +2070,8 @@ def test_python_parsed_module_never_falls_back_to_coarse_call(member_count): {"id": "caller_fetch", "label": "fetch()", "source_file": "caller.py"}, {"id": "helper", "label": "helper.py", "source_file": "helper.py"}] edges = [{"source": "caller", "target": "caller_fetch", "relation": "contains"}, - {"source": "caller", "target": "helper", "relation": "imports"}] + {"source": "caller", "target": "helper", "relation": "imports", + "_python_plain_import": True}] for number in range(member_count): nid = f"helper_get_{number}" nodes.append({"id": nid, "label": "get()", "source_file": "helper.py"}) @@ -2078,21 +2080,15 @@ def test_python_parsed_module_never_falls_back_to_coarse_call(member_count): assert not any(e["relation"] == "calls" for e in edges) -def test_python_external_module_call_survives_cache_and_context(tmp_path): +def test_python_external_module_call_survives_cache(tmp_path): source = tmp_path / "caller.py" source.write_text("import requests\n\ndef fetch():\n return requests.get('/data')\n") cold = extract([source], cache_root=tmp_path, root=tmp_path) warm = extract([source], cache_root=tmp_path, root=tmp_path) - context = extract( - [source], cache_root=tmp_path, root=tmp_path, - resolution_context_nodes=[{"id": "unrelated", "label": "unrelated.py", - "source_file": "unrelated.py"}], - resolution_context_edges=[], - ) def calls(result): return [e for e in result["edges"] if e["relation"] == "calls" and e["target"] == "requests"] - assert len(calls(cold)) == len(calls(warm)) == len(calls(context)) == 1 - assert calls(cold) == calls(warm) == calls(context) + assert len(calls(cold)) == len(calls(warm)) == 1 + assert calls(cold) == calls(warm) graph = build_from_json(cold) assert graph.nodes["requests"]["external"] is True assert graph.get_edge_data("caller_fetch", "requests")["relation"] == "calls" @@ -5410,3 +5406,41 @@ def test_3252_metadata_preservation(tmp_path): assert param_ref["source_location"] == "L2" assert node_by_id[param_ref["target"]]["label"] == "User" assert node_by_id[param_ref["target"]]["source_file"] == "models.py" + + +def test_python_external_calls_survive_real_incremental_context(tmp_path): + """A changed caller shares an external stub with an unchanged caller.""" + import subprocess + + (tmp_path / "ext_a.py").write_text("import requests\n\ndef fetch_a():\n return requests.get('/a')\n") + (tmp_path / "ext_b.py").write_text("import requests as rq\n\ndef fetch_b():\n return rq.post('/b')\n") + (tmp_path / "helper.py").write_text("def other():\n return 1\n") + (tmp_path / "local_caller.py").write_text("import helper\n\ndef fetch_local():\n return helper.get()\n") + env = {k: v for k, v in os.environ.items() + if k in {"PATH", "HOME", "TMPDIR", "LANG", "LC_ALL", "PYTHONDONTWRITEBYTECODE"}} + env["PYTHONPATH"] = str(Path(__file__).resolve().parents[1]) + + def run(): + result = subprocess.run( + [sys.executable, "-m", "graphify", "extract", str(tmp_path), "--code-only"], + cwd=tmp_path, env=env, capture_output=True, text=True, + ) + assert result.returncode == 0, result.stdout + result.stderr + graph = json.loads((tmp_path / "graphify-out" / "graph.json").read_text()) + calls = [e for e in graph["links"] if e.get("relation") == "calls"] + assert not any(e["source"] == "local_caller_fetch_local" for e in calls) + stub = next(n for n in graph["nodes"] if n["id"] == "requests") + assert stub["external"] is True and stub["source_file"] == "" + assert "_python_receiver_shadowed" not in json.dumps(graph) + assert "_python_plain_import" not in json.dumps(graph) + return result.stdout, sorted((e["source"], e["target"], e["source_location"]) for e in calls) + + full_stdout, full_calls = run() + assert "AST extraction on 4 code files" in full_stdout + assert full_calls == [("ext_a_fetch_a", "requests", "L4"), ("ext_b_fetch_b", "requests", "L4")] + (tmp_path / "ext_a.py").write_text("import requests\n\ndef fetch_a():\n # changed caller\n return requests.get('/a')\n") + (tmp_path / "local_caller.py").write_text("import helper\n\ndef fetch_local():\n # changed local consumer\n return helper.get()\n") + incremental_stdout, incremental_calls = run() + assert "incremental scan" in incremental_stdout + assert "AST extraction on 2 code files" in incremental_stdout + assert incremental_calls == [("ext_a_fetch_a", "requests", "L5"), ("ext_b_fetch_b", "requests", "L4")] diff --git a/tests/test_python_import_resolution.py b/tests/test_python_import_resolution.py index 49db038f1..f987f5cd2 100644 --- a/tests/test_python_import_resolution.py +++ b/tests/test_python_import_resolution.py @@ -41,6 +41,10 @@ def _has_edge(result: dict, source: str, target: str, relation: str) -> bool: ("import requests", "requests", "requests"), ("import requests as rq", "rq", "requests"), ("import requests as Requests", "Requests", "requests"), + ("import requests\nfrom requests import Session", "requests", "requests"), + ("from requests import Session; import requests", "requests", "requests"), + ("from other import value; import requests as rq", "rq", "requests"), + ("import requests as Requests\nfrom requests import Session", "Requests", "requests"), ("import os", "os", "os"), ("import pkg.sub as sub", "sub", "pkg_sub"), ]) @@ -61,7 +65,7 @@ def test_external_plain_import_member_call_targets_module( assert calls[0]["source"] == caller assert calls[0]["source_file"] == "caller.py" assert (calls[0]["confidence"], calls[0]["confidence_score"], calls[0]["context"], - calls[0]["source_location"], calls[0]["weight"]) == ("EXTRACTED", 1.0, "call", "L4", 1.0) + calls[0]["source_location"], calls[0]["weight"]) == ("EXTRACTED", 1.0, "call", f"L{len(import_line.splitlines()) + 3}", 1.0) assert not any(n.get("label") == "get()" for n in result["nodes"]) assert not any("_python_receiver_shadowed" in item for item in result["nodes"] + result["edges"]) @@ -89,6 +93,14 @@ def test_external_plain_import_member_call_targets_module( "def fetch():\n return (lambda *, requests=None: requests.get())()\n", "def fetch():\n return (lambda **requests: requests.get())()\n", "requests = object()\ndef fetch():\n return requests.get()\n", + "type requests = object\ndef fetch():\n return requests.get()\n", + "def fetch():\n type requests[T] = list[T]\n return requests.get()\n", + "def fetch[requests]():\n return requests.get()\n", + "import builtins\nbuiltins.exec('requests = object()')\ndef fetch():\n return requests.get()\n", + "import builtins as bi\nbi.eval('globals().update(requests=object())')\ndef fetch():\n return requests.get()\n", + "from builtins import exec as run\nrun('requests = object()')\ndef fetch():\n return requests.get()\n", + "import builtins\nrun = builtins.exec\nrun('requests = object()')\ndef fetch():\n return requests.get()\n", + "run = exec\nrun('requests = object()')\ndef fetch():\n return requests.get()\n", "def outer():\n requests = object()\n def fetch():\n return requests.get()\n", "if flag:\n import requests\ndef fetch():\n return requests.get()\n", "from elsewhere import *\ndef fetch():\n return requests.get()\n", @@ -108,6 +120,17 @@ def test_external_module_fallback_rejects_receiver_shadow(tmp_path: Path, body: assert not any(e["relation"] == "calls" and e["target"] == "requests" for e in result["edges"]) +@pytest.mark.parametrize("declaration", [ + "from requests import Session", + "from other import value as requests", +]) +def test_external_module_fallback_rejects_from_only_import(tmp_path: Path, declaration: str): + source = _write(tmp_path / "caller.py", f"{declaration}\ndef fetch():\n return requests.get()\n") + result = extract([source], cache_root=tmp_path) + assert not any(e["relation"] == "calls" and e["target"] in {"requests", "other"} + for e in result["edges"]) + + def test_overdeep_relative_import_is_unresolved_not_fatal(tmp_path: Path): source = _write( tmp_path / "pkg" / "mod.py",