diff --git a/graphify/extractors/powershell.py b/graphify/extractors/powershell.py index d570475ae..43c0aaae6 100644 --- a/graphify/extractors/powershell.py +++ b/graphify/extractors/powershell.py @@ -30,7 +30,10 @@ def extract_powershell(path: Path) -> dict: nodes: list[dict] = [] edges: list[dict] = [] seen_ids: set[str] = set() - function_bodies: list[tuple[str, Any]] = [] + function_bodies: list[tuple[str, Any, str | None]] = [] + class_ids_by_name: dict[str, list[str]] = {} + class_names_by_id: dict[str, str] = {} + methods_by_class: dict[str, dict[str, str]] = {} def add_node(nid: str, label: str, line: int) -> None: if nid not in seen_ids: @@ -120,7 +123,7 @@ def extract_powershell(path: Path) -> dict: add_edge(file_nid, func_nid, "contains", line) body = _find_script_block_body(node) if body: - function_bodies.append((func_nid, body)) + function_bodies.append((func_nid, body, None)) # Also walk the body during the main pass so that # Import-Module / dot-source inside functions emit # file-level imports_from edges (#1331). @@ -165,6 +168,9 @@ def extract_powershell(path: Path) -> dict: class_name = _read_text(name_node, source) line = node.start_point[0] + 1 class_nid = _make_id(stem, class_name) + # Count declarations, including duplicates that share a node id. + class_ids_by_name.setdefault(class_name.lower(), []).append(class_nid) + class_names_by_id[class_nid] = class_name.lower() add_node(class_nid, class_name, line) add_edge(file_nid, class_nid, "contains", line) # Base type(s) after ':'. PowerShell has no syntactic base vs @@ -207,6 +213,7 @@ def extract_powershell(path: Path) -> dict: method_nid = _make_id(parent_class_nid, method_name) add_node(method_nid, f".{method_name}()", line) add_edge(parent_class_nid, method_nid, "method", line) + methods_by_class.setdefault(parent_class_nid, {})[method_name.lower()] = method_nid else: method_nid = _make_id(stem, method_name) add_node(method_nid, f"{method_name}()", line) @@ -239,7 +246,7 @@ def extract_powershell(path: Path) -> dict: p_line, context="parameter_type") body = _find_script_block_body(node) if body: - function_bodies.append((method_nid, body)) + function_bodies.append((method_nid, body, parent_class_nid)) return if t == "command": @@ -315,17 +322,10 @@ def extract_powershell(path: Path) -> dict: walk(root) label_to_nid = {n["label"].strip("()").lstrip(".").lower(): n["id"] for n in nodes} - # Member calls (`$this.M()`, `[C]::M()`) must resolve only to a method, never - # to a free function that happens to share the name. Methods are exactly the - # nodes whose label is dot-prefixed (`.M()`), so key them separately (#3992). - method_label_to_nid = { - n["label"].strip("()").lstrip(".").lower(): n["id"] - for n in nodes if n["label"].startswith(".") - } seen_call_pairs: set[tuple[str, str]] = set() raw_calls: list[dict] = [] - def walk_calls(node, caller_nid: str) -> None: + def walk_calls(node, caller_nid: str, parent_class_nid: str | None) -> None: if node.type in ("function_statement", "class_statement"): return if node.type == "command": @@ -360,7 +360,26 @@ def extract_powershell(path: Path) -> dict: member_node = next((c for c in node.children if c.type == "member_name"), None) if member_node is not None: method_name = _read_text(member_node, source) - tgt_nid = method_label_to_nid.get(method_name.lower()) + receiver_class = None + receiver = node.children[0] if node.children else None + operator = node.children[1].type if len(node.children) > 1 else None + # Only a direct lexical self or unique literal local type proves + # ownership. Value, expression and dynamic receivers stay raw. + if receiver is not None and receiver.type == "variable" and operator == ".": + if _read_text(receiver, source).lower() == "$this" and parent_class_nid: + name = class_names_by_id[parent_class_nid] + if len(class_ids_by_name[name]) == 1: + receiver_class = parent_class_nid + elif receiver is not None and receiver.type == "type_literal" and operator == "::": + # Match the whole literal: a partial parse can hide type qualifiers. + type_name = _read_text(receiver, source)[1:-1].strip() if not receiver.has_error else "" + class_ids = class_ids_by_name.get(type_name.lower(), []) if type_name else [] + if len(class_ids) == 1: + receiver_class = class_ids[0] + literal_member = (len(member_node.named_children) == 1 + and member_node.named_children[0].type == "simple_name") + tgt_nid = (methods_by_class.get(receiver_class, {}).get(method_name.lower()) + if receiver_class and literal_member else None) if tgt_nid and tgt_nid != caller_nid: pair = (caller_nid, tgt_nid) if pair not in seen_call_pairs: @@ -377,10 +396,10 @@ def extract_powershell(path: Path) -> dict: "source_location": f"L{node.start_point[0] + 1}", }) for child in node.children: - walk_calls(child, caller_nid) + walk_calls(child, caller_nid, parent_class_nid) - for caller_nid, body_node in function_bodies: - walk_calls(body_node, caller_nid) + for caller_nid, body_node, parent_class_nid in function_bodies: + walk_calls(body_node, caller_nid, parent_class_nid) clean_edges = [e for e in edges if e["source"] in seen_ids and (e["target"] in seen_ids or e["relation"] in ("imports_from", "imports"))] diff --git a/tests/test_languages.py b/tests/test_languages.py index 52b5fc120..a0e2f67cb 100644 --- a/tests/test_languages.py +++ b/tests/test_languages.py @@ -2740,6 +2740,131 @@ def test_powershell_command_call_still_resolves(tmp_path): assert ("Run", "Helper") in _edge_labels(r, "calls") +@pytest.mark.parametrize("reverse", [False, True]) +def test_powershell_member_calls_keep_receiver_class(tmp_path, reverse): + classes = [ + f"class {name} {{\n" + " [int] Pick() { return 1 }\n" + " static [int] Make() { return 1 }\n" + " [int] Run() { $THIS.pIcK(); return 1 }\n" + "}\n" + for name in ("Alpha", "Beta") + ] + f = tmp_path / "owners.ps1" + f.write_text("".join(reversed(classes) if reverse else classes) + + "function Start { [aLpHa]::mAkE(); [Beta]::Make() }\n") + r = extract_powershell(f) + assert "error" not in r + classes_by_id = {n["id"]: n["label"] for n in r["nodes"] if n["label"] in ("Alpha", "Beta")} + labels = {n["id"]: n["label"] for n in r["nodes"]} + methods = {(classes_by_id[e["source"]], labels[e["target"]]): e["target"] + for e in r["edges"] if e["relation"] == "method"} + calls = {(e["source"], e["target"]) for e in r["edges"] if e["relation"] == "calls"} + start = next(n["id"] for n in r["nodes"] if n["label"] == "Start()") + assert calls == { + (methods[(name, ".Run()")], methods[(name, ".Pick()")]) for name in ("Alpha", "Beta") + } | {(start, methods[(name, ".Make()")]) for name in ("Alpha", "Beta")} + + +@pytest.mark.parametrize("call,callee", [ + ("$unknown.Pick()", "Pick"), + ("$typed.Pick()", "Pick"), + ("$local.Pick()", "Pick"), + ("$this.child.Pick()", "Pick"), + ("([Alpha]$this).Pick()", "Pick"), + ("$this.$dynamic()", "$dynamic"), + ("[External]::Pick()", "Pick"), + ("[Alpha+Nested]::Pick()", "Pick"), + ("[Alpha, MyAssembly]::Pick()", "Pick"), + ("[Beta]::Pick()", "Pick"), +]) +def test_powershell_uncertain_receivers_remain_raw(tmp_path, call, callee): + f = tmp_path / "unknown.ps1" + f.write_text( + "function Pick { return 3 }\n" + "class Alpha {\n" + " [int] Pick() { return 1 }\n" + " [int] Run([Alpha]$typed) {\n" + " $local = [Alpha]::new()\n" + f" {call}\n" + " return 1\n" + " }\n" + "}\n" + "class Beta { [int] Other() { return 1 } }\n" + ) + r = extract_powershell(f) + assert "error" not in r + assert not [e for e in r["edges"] if e["relation"] == "calls"] + picked = [c for c in r["raw_calls"] if c["callee"] == callee] + assert len(picked) == 1 + assert picked[0]["is_member_call"] is True + assert picked[0]["source_file"] == str(f) + assert picked[0]["source_location"].startswith("L") + + +@pytest.mark.parametrize("receiver", ["[Alpha]", "[aLpHa]", "[ Alpha ]"]) +def test_powershell_complete_static_type_literal_resolves(tmp_path, receiver): + f = tmp_path / "literal.ps1" + f.write_text( + "class Alpha { static [int] Pick() { return 1 } }\n" + f"function Start {{ {receiver}::Pick() }}\n" + ) + r = extract_powershell(f) + assert "error" not in r + caller = next(n["id"] for n in r["nodes"] if n["label"] == "Start()") + target = next(n["id"] for n in r["nodes"] if n["label"] == ".Pick()") + calls = [e for e in r["edges"] if e["relation"] == "calls"] + assert [(e["source"], e["target"]) for e in calls] == [(caller, target)] + assert not [c for c in r["raw_calls"] if c["callee"] == "Pick"] + + +@pytest.mark.parametrize("second_name", ["Alpha", "alpha"]) +def test_powershell_duplicate_class_names_do_not_guess(tmp_path, second_name): + f = tmp_path / "duplicate.ps1" + f.write_text( + "class Alpha { [int] Pick() { return 1 } [int] Run() { return $this.Pick() } }\n" + f"class {second_name} {{ [int] Pick() {{ return 2 }} }}\n" + "function Start { [Alpha]::Pick() }\n" + ) + r = extract_powershell(f) + assert "error" not in r + assert not [e for e in r["edges"] if e["relation"] == "calls"] + assert len([c for c in r["raw_calls"] if c["callee"] == "Pick" and c["is_member_call"]]) == 2 + + +def test_powershell_property_reads_and_nested_calls(tmp_path): + f = tmp_path / "nested.ps1" + f.write_text( + "class Alpha {\n" + " static [int] Make() { return 1 }\n" + " [int] Run() { $this.Make; [Alpha]::Make; $unknown.Pick([Alpha]::Make()); return 1 }\n" + "}\n" + ) + r = extract_powershell(f) + assert _edge_labels(r, "calls") == {("Run", "Make")} + assert [c["callee"] for c in r["raw_calls"]] == ["Pick"] + + +def test_powershell_unknown_member_stays_unresolved_in_cold_and_warm_corpus(tmp_path): + from graphify.extract import extract + f = tmp_path / "corpus.ps1" + f.write_text( + "function Pick { return 3 }\n" + "class Alpha { [int] Pick() { return 1 } [int] Run() { return $this.Pick() } }\n" + "class Beta { [int] Pick() { return 2 } }\n" + "function Start { $unknown.Pick() }\n" + ) + cold = extract([f], cache_root=tmp_path / "cache", root=tmp_path, parallel=False) + warm = extract([f], cache_root=tmp_path / "cache", root=tmp_path, parallel=False) + for r in (cold, warm): + ids = {n["label"]: n["id"] for n in r["nodes"] if n["label"] in ("Alpha", "Beta", "Start()")} + methods = {(e["source"], next(n["label"] for n in r["nodes"] if n["id"] == e["target"])): e["target"] + for e in r["edges"] if e["relation"] == "method"} + assert {(e["source"], e["target"]) for e in r["edges"] if e["relation"] == "calls"} == { + (methods[(ids["Alpha"], ".Run()")], methods[(ids["Alpha"], ".Pick()")]) + } + + def test_powershell_class_base_type_emits_inherits_edge(): # `class Circle : Shape` — the base type after ':' was previously dropped # because the handler only read the first simple_name (the class name).