From 80e98090866bf4745f71cbb780ac5fe87cc752c3 Mon Sep 17 00:00:00 2001 From: mochan Date: Wed, 10 Jun 2026 18:40:09 +0530 Subject: [PATCH] feat(graph): Python intra-procedural data edges (DATA_ASSIGN/ARG/RETURN) + PARAMETER nodes Co-Authored-By: Claude Sonnet 4.6 --- codegraph/graph/schema.py | 4 + codegraph/parsers/python.py | 616 ++++++++++++++++++++++++++++++++ tests/test_data_edges_python.py | 435 ++++++++++++++++++++++ 3 files changed, 1055 insertions(+) create mode 100644 tests/test_data_edges_python.py diff --git a/codegraph/graph/schema.py b/codegraph/graph/schema.py index fe6efda..391fbff 100644 --- a/codegraph/graph/schema.py +++ b/codegraph/graph/schema.py @@ -37,6 +37,10 @@ class EdgeKind(str, Enum): READS_FROM = "READS_FROM" # function → SQLAlchemy model on read (DF1) WRITES_TO = "WRITES_TO" # function → SQLAlchemy model on write (DF1) FETCH_CALL = "FETCH_CALL" # frontend call site → URL string (DF2, fetch/axios) + # Intra-procedural data-flow edges (C1 — Python body analysis). + DATA_ASSIGN = "DATA_ASSIGN" # source var/param → target variable (assignment) + DATA_ARG = "DATA_ARG" # variable/param → call-site argument sentinel + DATA_RETURN = "DATA_RETURN" # variable/param → function return sentinel class Node(BaseModel): diff --git a/codegraph/parsers/python.py b/codegraph/parsers/python.py index 4117637..5c8c5f5 100644 --- a/codegraph/parsers/python.py +++ b/codegraph/parsers/python.py @@ -906,6 +906,12 @@ class PythonExtractor(ExtractorBase): # ``_is_entry_point``. extra_entry_point_decorators: tuple[str, ...] = () + # Kill-switch for intra-procedural data-flow edge emission (C1). + # Set to False to disable DATA_ASSIGN / DATA_ARG / DATA_RETURN edges + # and PARAMETER node emission entirely. Attribute hook for future config + # file plumbing — the builder's _apply_config_to_extractors can flip it. + emit_data_edges: bool = True + def parse_file( self, path: Path, repo_root: Path ) -> tuple[list[Node], list[Edge]]: @@ -1247,12 +1253,29 @@ def _handle_function( nodes, edges, ) + # C1 — PARAMETER nodes + intra-procedural data edges. + # Emit one PARAMETER node per logical parameter (already filtered by + # skip_self_or_cls in metadata["params"]). We also collect the raw + # params_node here so we can correlate line numbers. + param_scope: dict[str, str] = {} # name → node-id for data-edge seeding + if self.emit_data_edges and params is not None: + param_scope = self._emit_param_nodes( + params, src, rel, qualname, func_id, + kind == NodeKind.METHOD, node.start_point[0] + 1, + nodes, edges, + ) + if body is not None: self._collect_calls(body, rel, func_id, src, edges) # DF1 — SQLAlchemy READS_FROM / WRITES_TO. Walk the body for # ORM session calls; emits ``unresolved::Model`` edges that # the post-build resolver rewrites to real CLASS ids. self._collect_sql_io(body, rel, func_id, src, edges) + # C1 — intra-procedural data-flow edges (DATA_ASSIGN/ARG/RETURN). + if self.emit_data_edges: + self._collect_data_edges( + body, rel, func_id, qualname, src, param_scope, nodes, edges, + ) # Visit nested defs so their bodies and calls are not lost. # The innermost named function owns its calls — that mirrors # the runtime attribution and matches what users expect when @@ -1262,6 +1285,599 @@ def _handle_function( src, nodes, edges, ) + def _emit_param_nodes( + self, + params_node: tree_sitter.Node, + src: bytes, + rel: str, + func_qualname: str, + func_id: str, + skip_self_or_cls: bool, + func_line: int, + nodes: list[Node], + edges: list[Edge], + ) -> dict[str, str]: + """Emit a PARAMETER node and PARAM_OF edge for each function parameter. + + Returns a scope seed dict ``{param_name: param_node_id}`` for use by + ``_collect_data_edges``. + + Limitations (documented): + - Variadic params (*args, **kwargs) are included with their prefixed + names and cannot be individually resolved in the data-flow scope. + - self/cls is skipped for methods, consistent with DF0 behaviour. + """ + scope: dict[str, str] = {} + first_seen = False + index = 0 + for child in params_node.children: + if not child.is_named: + continue + # Determine name, type, line from the parameter AST node. + param_name: str | None = None + param_type: str | None = None + param_line = child.start_point[0] + 1 + if child.type == "identifier": + param_name = node_text(child, src) + elif child.type == "typed_parameter": + name_n = next( + (c for c in child.children if c.type == "identifier"), None + ) + type_n = next( + (c for c in child.children if c.type == "type"), None + ) + if name_n is not None: + param_name = node_text(name_n, src) + param_type = node_text(type_n, src) if type_n else None + elif child.type == "default_parameter": + name_n = child.child_by_field_name("name") + if name_n is not None: + param_name = node_text(name_n, src) + elif child.type == "typed_default_parameter": + name_n = child.child_by_field_name("name") + type_n = child.child_by_field_name("type") + if name_n is not None: + param_name = node_text(name_n, src) + param_type = node_text(type_n, src) if type_n else None + elif child.type == "list_splat_pattern": + inner = next( + (c for c in child.children if c.type == "identifier"), None + ) + if inner is not None: + param_name = f"*{node_text(inner, src)}" + elif child.type == "dictionary_splat_pattern": + inner = next( + (c for c in child.children if c.type == "identifier"), None + ) + if inner is not None: + param_name = f"**{node_text(inner, src)}" + if param_name is None: + continue + if ( + skip_self_or_cls + and not first_seen + and param_name in ("self", "cls") + ): + first_seen = True + continue + first_seen = True + param_qualname = f"{func_qualname}." + param_id = make_node_id(NodeKind.PARAMETER, param_qualname, rel) + nodes.append(Node( + id=param_id, + kind=NodeKind.PARAMETER, + name=param_name, + qualname=param_qualname, + file=rel, + line_start=param_line, + line_end=param_line, + language="python", + metadata={ + "index": index, + "name": param_name, + "type": param_type, + }, + )) + edges.append(Edge( + src=param_id, + dst=func_id, + kind=EdgeKind.PARAM_OF, + file=rel, + line=func_line, + )) + # Seed the scope with the bare name (strip variadic prefix for + # *args/**kwargs so callers referencing ``args`` in the body work). + bare_name = param_name.lstrip("*") + if bare_name: + scope[bare_name] = param_id + index += 1 + return scope + + def _collect_data_edges( + self, + body: tree_sitter.Node, + rel: str, + func_id: str, + func_qualname: str, + src: bytes, + param_scope: dict[str, str], + nodes: list[Node], + edges: list[Edge], + ) -> None: + """Walk a function body and emit intra-procedural data-flow edges. + + Emits DATA_ASSIGN, DATA_ARG, and DATA_RETURN edges for simple scalar + variable flow within a single function. The scope maps variable names + to their most-recent definition node-id (a PARAMETER or VARIABLE node). + + **Design decisions / limitations (intentional simplifications)**: + + * Single-level shadowing by line: ``x = a; x = b`` emits two distinct + VARIABLE nodes qualified by line, so shadow chains are traceable. + * No attribute tracking: ``obj.field = x`` is not captured. + * No tuple unpacking: ``a, b = f()`` — the first target only is taken + when the LHS is a tuple; the others are skipped. + * Comprehension scopes: treated as part of the surrounding function + scope (no inner binding isolation). + * Calls in RHS: ``x = foo(a)`` emits DATA_ASSIGN from the synthetic + sentinel ``unresolved::ret::foo`` → VARIABLE x, metadata + ``{"callee": "foo"}``. Cross-file binding of the sentinel to the + callee's return node happens in a later PR. + * Only plain identifiers in RHS are tracked — binary ops, subscripts, + and attribute access on in-scope names do NOT propagate. This keeps + the implementation correct (no false positives) at the cost of + missing some true positives. + """ + # Mutable scope: maps current name → most recent PARAMETER/VARIABLE id. + scope: dict[str, str] = dict(param_scope) + # Track the return-sentinel node-id lazily (emit once per function). + return_node_id: str | None = None + + # Walk statements in document order to maintain correct scope. + # We use a recursive helper to preserve forward-order processing + # while descending into control-flow blocks. + return_node_id = self._walk_data_stmts( + list(body.children), rel, func_id, func_qualname, src, + scope, return_node_id, nodes, edges, + ) + + def _walk_data_stmts( + self, + stmts: list[tree_sitter.Node], + rel: str, + func_id: str, + func_qualname: str, + src: bytes, + scope: dict[str, str], + return_node_id: str | None, + nodes: list[Node], + edges: list[Edge], + ) -> str | None: + """Walk *stmts* in document order, updating scope and emitting data edges. + + Returns the (possibly newly created) return-sentinel node-id. + """ + for stmt in stmts: + # --- assignment: x = expr (and annotated: x: T = expr) ------ + if stmt.type == "expression_statement": + inner = next((c for c in stmt.children if c.is_named), None) + if inner is not None and inner.type in ( + "assignment", "augmented_assignment", + ): + self._handle_data_assign( + inner, rel, func_id, func_qualname, src, scope, + nodes, edges, + ) + elif inner is not None and inner.type == "call": + # Standalone call statement: ``g(x)`` / ``g(user=x)``. + # Emit DATA_ARG edges for in-scope identifier arguments. + func_child = inner.child_by_field_name("function") + callee_text = ( + node_text(func_child, src) if func_child else "" + ) + arg_list = inner.child_by_field_name("arguments") + if arg_list is not None: + self._emit_arg_edges( + arg_list, callee_text, rel, src, scope, edges, + inner.start_point[0] + 1, + ) + + # --- return expr ----------------------------------------------- + elif stmt.type == "return_statement": + return_node_id = self._handle_data_return( + stmt, rel, func_qualname, src, scope, + return_node_id, nodes, edges, + ) + + # --- for target: ``for x in iterable:`` ------------------------ + # The loop variable is a definition fed by the iterable — + # conservative element-of propagation (taint on the collection + # taints the element). Critical for source patterns like + # ``for url in scraped_urls: fetch(url)``. + elif stmt.type == "for_statement": + self._handle_for_target( + stmt, rel, func_id, func_qualname, src, scope, + nodes, edges, + ) + body = stmt.child_by_field_name("body") + if body is not None: + return_node_id = self._walk_data_stmts( + list(body.children), rel, func_id, func_qualname, + src, scope, return_node_id, nodes, edges, + ) + + # Descend into control-flow blocks but stop at nested defs. + # We preserve document order by recursing rather than stacking. + elif stmt.type not in ( + "function_definition", "class_definition", "decorator", + ): + for child in stmt.children: + if child.type == "block": + return_node_id = self._walk_data_stmts( + list(child.children), rel, func_id, func_qualname, + src, scope, return_node_id, nodes, edges, + ) + elif child.type in ( + "if_statement", "for_statement", "while_statement", + "with_statement", "try_statement", + "expression_statement", "return_statement", + ): + return_node_id = self._walk_data_stmts( + [child], rel, func_id, func_qualname, + src, scope, return_node_id, nodes, edges, + ) + return return_node_id + + def _handle_data_assign( + self, + assign_node: tree_sitter.Node, + rel: str, + func_id: str, + func_qualname: str, + src: bytes, + scope: dict[str, str], + nodes: list[Node], + edges: list[Edge], + ) -> None: + """Process one assignment or augmented-assignment statement.""" + line = assign_node.start_point[0] + 1 + + if assign_node.type == "augmented_assignment": + # ``x += rhs`` — treat as rhs → x, but x must already be in scope + # (we don't re-define it since aug-assign modifies in place). + lhs_node = assign_node.child_by_field_name("left") + rhs_node = assign_node.child_by_field_name("right") + if lhs_node is None or rhs_node is None: + return + lhs_name = ( + node_text(lhs_node, src) + if lhs_node.type == "identifier" + else None + ) + if lhs_name is None or lhs_name not in scope: + return + var_id = scope[lhs_name] + self._emit_rhs_edges( + rhs_node, rel, var_id, func_qualname, src, scope, edges, line, + ) + return + + # Standard assignment (covers plain ``x = y`` and annotated ``x: T = y``). + # LHS: use the first simple identifier target. + lhs_node = assign_node.child_by_field_name("left") + rhs_node = assign_node.child_by_field_name("right") + # Also handle annotated assignment where the field name is "name" + # (tree-sitter grammar variation). + if lhs_node is None: + lhs_node = next( + ( + c for c in assign_node.children + if c.is_named and c.type == "identifier" + ), + None, + ) + if rhs_node is None or lhs_node is None: + return + + # Tuple unpacking — take first simple identifier only. + if lhs_node.type in ("tuple_pattern", "tuple"): + lhs_node = next( + (c for c in lhs_node.children if c.is_named and c.type == "identifier"), + None, + ) + if lhs_node is None or lhs_node.type != "identifier": + return + + lhs_name = node_text(lhs_node, src) + if not lhs_name: + return + + # Create a new VARIABLE node (line-qualified so shadowing is distinct). + var_qualname = f"{func_qualname}." + var_id = make_node_id(NodeKind.VARIABLE, var_qualname, rel) + nodes.append(Node( + id=var_id, + kind=NodeKind.VARIABLE, + name=lhs_name, + qualname=var_qualname, + file=rel, + line_start=line, + line_end=line, + language="python", + metadata={"assigned_in": func_qualname}, + )) + edges.append(Edge( + src=var_id, + dst=func_id, + kind=EdgeKind.DEFINED_IN, + file=rel, + line=line, + )) + + # Update scope BEFORE emitting RHS edges so self-referential + # assignments (``x = x + 1``) use the *old* binding for the RHS. + old_id = scope.get(lhs_name) + scope[lhs_name] = var_id + + self._emit_rhs_edges( + rhs_node, rel, var_id, func_qualname, src, + # Use scope snapshot with old binding for the RHS reads. + {**scope, lhs_name: old_id} if old_id else scope, + edges, line, + ) + + def _handle_for_target( + self, + for_node: tree_sitter.Node, + rel: str, + func_id: str, + func_qualname: str, + src: bytes, + scope: dict[str, str], + nodes: list[Node], + edges: list[Edge], + ) -> None: + """Register a for-loop target as a scoped VARIABLE fed by its iterable. + + ``for x in items`` emits DATA_ASSIGN edges from in-scope identifiers + inside ``items`` to a new VARIABLE node for ``x`` (element-of + propagation). Tuple targets take the first simple identifier only, + mirroring :meth:`_handle_data_assign`. + """ + left = for_node.child_by_field_name("left") + right = for_node.child_by_field_name("right") + if left is not None and left.type in ("tuple_pattern", "tuple", "pattern_list"): + left = next( + (c for c in left.children if c.is_named and c.type == "identifier"), + None, + ) + if left is None or left.type != "identifier": + return + name = node_text(left, src) + if not name: + return + line = left.start_point[0] + 1 + var_qualname = f"{func_qualname}." + var_id = make_node_id(NodeKind.VARIABLE, var_qualname, rel) + nodes.append(Node( + id=var_id, + kind=NodeKind.VARIABLE, + name=name, + qualname=var_qualname, + file=rel, + line_start=line, + line_end=line, + language="python", + metadata={"assigned_in": func_qualname, "loop_target": True}, + )) + edges.append(Edge( + src=var_id, + dst=func_id, + kind=EdgeKind.DEFINED_IN, + file=rel, + line=line, + )) + if right is not None: + self._emit_rhs_edges( + right, rel, var_id, func_qualname, src, scope, edges, line, + ) + scope[name] = var_id + + def _emit_rhs_edges( + self, + rhs_node: tree_sitter.Node, + rel: str, + dst_id: str, + func_qualname: str, + src: bytes, + scope: dict[str, str], + edges: list[Edge], + line: int, + ) -> None: + """Emit DATA_ASSIGN edges from identifiers/call-results in *rhs_node*. + + Walks the RHS expression tree collecting: + * Plain identifier → DATA_ASSIGN if name is in scope. + * ``call`` node → DATA_ASSIGN from the sentinel + ``unresolved::ret::``, and also DATA_ARG edges for each + plain-identifier argument that resolves in scope. + """ + # BFS over the RHS to collect all identifiers and calls. + rhs_stack = [rhs_node] + while rhs_stack: + node = rhs_stack.pop() + if node.type == "identifier": + name = node_text(node, src) + if name and name in scope and scope[name] is not None: + edges.append(Edge( + src=scope[name], + dst=dst_id, + kind=EdgeKind.DATA_ASSIGN, + file=rel, + line=line, + )) + elif node.type == "call": + func_child = node.child_by_field_name("function") + callee_text = node_text(func_child, src) if func_child else "" + # DATA_ASSIGN from return-value sentinel. + edges.append(Edge( + src=f"unresolved::ret::{callee_text}", + dst=dst_id, + kind=EdgeKind.DATA_ASSIGN, + file=rel, + line=line, + metadata={"callee": callee_text}, + )) + # DATA_ARG edges for in-scope arguments. + arg_list = node.child_by_field_name("arguments") + if arg_list is not None: + self._emit_arg_edges( + arg_list, callee_text, rel, src, scope, edges, line, + ) + # Do NOT recurse into nested calls via rhs_stack here — the + # call itself already represents the value, and recursing + # would incorrectly attribute identifier nodes inside the + # argument list as direct sources of dst_id. + continue + else: + rhs_stack.extend(c for c in node.children if c.is_named) + + def _emit_arg_edges( + self, + arg_list: tree_sitter.Node, + callee_text: str, + rel: str, + src: bytes, + scope: dict[str, str], + edges: list[Edge], + line: int, + ) -> None: + """Emit DATA_ARG edges for plain-identifier arguments in a call. + + Cross-file binding of the sentinel dst to the callee's PARAMETER node + is deferred to a later PR. Here we emit unresolved sentinels of the + form ``unresolved::arg::::``. + + Limitations: only plain identifier arguments produce edges; complex + expressions (``f(a + b)``, ``f(obj.attr)``) are silently skipped. + """ + pos = 0 + for child in arg_list.children: + if not child.is_named: + continue + if child.type == "keyword_argument": + kw_name_n = child.child_by_field_name("name") + kw_val_n = child.child_by_field_name("value") + if ( + kw_name_n is not None + and kw_val_n is not None + and kw_val_n.type == "identifier" + ): + arg_name = node_text(kw_val_n, src) + kw = node_text(kw_name_n, src) + if arg_name and arg_name in scope and scope[arg_name] is not None: + edges.append(Edge( + src=scope[arg_name], + dst=f"unresolved::arg::{callee_text}::{kw}", + kind=EdgeKind.DATA_ARG, + file=rel, + line=line, + metadata={ + "callee": callee_text, + "kwarg": kw, + "call_line": line, + }, + )) + elif child.type == "identifier": + arg_name = node_text(child, src) + if arg_name and arg_name in scope and scope[arg_name] is not None: + edges.append(Edge( + src=scope[arg_name], + dst=f"unresolved::arg::{callee_text}::{pos}", + kind=EdgeKind.DATA_ARG, + file=rel, + line=line, + metadata={ + "callee": callee_text, + "position": pos, + "call_line": line, + }, + )) + pos += 1 + else: + pos += 1 + + def _handle_data_return( + self, + return_node: tree_sitter.Node, + rel: str, + func_qualname: str, + src: bytes, + scope: dict[str, str], + return_node_id: str | None, + nodes: list[Node], + edges: list[Edge], + ) -> str | None: + """Emit DATA_RETURN edges from in-scope identifiers in a return expr. + + The synthetic return sentinel node (kind VARIABLE, + qualname ``.``) is created lazily on first return + encountered and reused for all subsequent returns in the same function. + + Returns the sentinel's node-id (possibly newly created). + """ + # Find the expression being returned. + expr_node = next( + (c for c in return_node.children if c.is_named and c.type != "return"), + None, + ) + if expr_node is None: + return return_node_id + + # Collect plain identifiers in the return expression. + id_names: list[str] = [] + expr_stack = [expr_node] + while expr_stack: + n = expr_stack.pop() + if n.type == "identifier": + id_names.append(node_text(n, src)) + else: + expr_stack.extend(c for c in n.children if c.is_named) + + sources = [ + scope[nm] for nm in id_names + if nm in scope and scope[nm] is not None + ] + if not sources: + return return_node_id + + # Create return sentinel node lazily. + if return_node_id is None: + ret_qualname = f"{func_qualname}." + return_node_id = make_node_id(NodeKind.VARIABLE, ret_qualname, rel) + nodes.append(Node( + id=return_node_id, + kind=NodeKind.VARIABLE, + name="", + qualname=ret_qualname, + file=rel, + line_start=return_node.start_point[0] + 1, + line_end=return_node.start_point[0] + 1, + language="python", + metadata={"synthetic_kind": "RETURN", "func": func_qualname}, + )) + + line = return_node.start_point[0] + 1 + for src_id in sources: + edges.append(Edge( + src=src_id, + dst=return_node_id, + kind=EdgeKind.DATA_RETURN, + file=rel, + line=line, + )) + return return_node_id + def _visit_nested_defs( self, block: tree_sitter.Node, diff --git a/tests/test_data_edges_python.py b/tests/test_data_edges_python.py new file mode 100644 index 0000000..35e3db3 --- /dev/null +++ b/tests/test_data_edges_python.py @@ -0,0 +1,435 @@ +"""C1: Python intra-procedural data-flow edges + PARAMETER nodes. + +Tests cover: +* PARAMETER nodes emitted for each function param with PARAM_OF edges. +* DATA_ASSIGN: parameter → variable, variable → variable. +* DATA_ARG: in-scope variable to call-site sentinel (positional and keyword). +* DATA_RETURN: in-scope variable to function return sentinel. +* Call result assignment: DATA_ASSIGN from unresolved::ret::. +* Shadowing: two VARIABLE nodes with distinct line-qualified qualnames. +* Nested def: inner function variables do NOT leak into outer scope. +* Kill-switch: emit_data_edges=False suppresses all C1 edges/nodes. +""" +from __future__ import annotations + +import shutil +import time +from pathlib import Path + +import pytest + +from codegraph.graph.schema import Edge, EdgeKind, Node, NodeKind +from codegraph.parsers.python import PythonExtractor + +# --------------------------------------------------------------------------- +# Helpers +# --------------------------------------------------------------------------- + +def _run(tmp_path: Path, src: str, filename: str = "sample.py") -> tuple[list[Node], list[Edge]]: + target = tmp_path / filename + target.write_text(src) + return PythonExtractor().parse_file(target, tmp_path) + + +def _data_edges(edges: list[Edge], kind: EdgeKind) -> list[Edge]: + return [e for e in edges if e.kind == kind] + + +def _nodes_of_kind(nodes: list[Node], kind: NodeKind) -> list[Node]: + return [n for n in nodes if n.kind == kind] + + +def _param_node(nodes: list[Node], param_name: str) -> Node | None: + return next( + (n for n in nodes if n.kind == NodeKind.PARAMETER and n.name == param_name), + None, + ) + + +def _var_node(nodes: list[Node], var_name: str) -> Node | None: + """Return first VARIABLE node with the given name (not a sentinel).""" + return next( + ( + n for n in nodes + if n.kind == NodeKind.VARIABLE + and n.name == var_name + and not n.qualname.endswith(".") + and "synthetic_kind" not in n.metadata + ), + None, + ) + + +def _return_sentinel(nodes: list[Node], func_qualname: str) -> Node | None: + expected_qualname = f"{func_qualname}." + return next( + (n for n in nodes if n.qualname == expected_qualname), + None, + ) + + +# --------------------------------------------------------------------------- +# PARAMETER node emission +# --------------------------------------------------------------------------- + + +def test_param_nodes_emitted(tmp_path: Path) -> None: + """Each function parameter gets a PARAMETER node.""" + src = "def f(a, b, c):\n pass\n" + nodes, _ = _run(tmp_path, src) + params = _nodes_of_kind(nodes, NodeKind.PARAMETER) + assert len(params) == 3 + names = {p.name for p in params} + assert names == {"a", "b", "c"} + + +def test_param_node_qualnames(tmp_path: Path) -> None: + """PARAMETER qualnames follow the . convention.""" + src = "def f(x):\n pass\n" + nodes, _ = _run(tmp_path, src) + p = _param_node(nodes, "x") + assert p is not None + assert p.qualname == "sample.f." + + +def test_param_node_metadata(tmp_path: Path) -> None: + """PARAMETER node metadata includes index, name, and type.""" + src = "def f(a, b: int):\n pass\n" + nodes, _ = _run(tmp_path, src) + a = _param_node(nodes, "a") + b = _param_node(nodes, "b") + assert a is not None and a.metadata["index"] == 0 and a.metadata["type"] is None + assert b is not None and b.metadata["index"] == 1 and b.metadata["type"] == "int" + + +def test_param_of_edges_emitted(tmp_path: Path) -> None: + """Each PARAMETER node has a PARAM_OF edge pointing to the function.""" + src = "def f(a, b):\n pass\n" + nodes, edges = _run(tmp_path, src) + param_of = _data_edges(edges, EdgeKind.PARAM_OF) + assert len(param_of) == 2 + func_node = next(n for n in nodes if n.kind == NodeKind.FUNCTION) + for e in param_of: + assert e.dst == func_node.id + + +def test_param_count_matches_signature(tmp_path: Path) -> None: + """PARAMETER node count equals the number of parameters (excluding self).""" + src = "class C:\n def m(self, x: str, y: int = 0):\n pass\n" + nodes, _ = _run(tmp_path, src) + params = _nodes_of_kind(nodes, NodeKind.PARAMETER) + # self is skipped, so only x and y + assert len(params) == 2 + names = {p.name for p in params} + assert names == {"x", "y"} + + +def test_variadic_param_nodes(tmp_path: Path) -> None: + """Variadic *args / **kwargs also get PARAMETER nodes.""" + src = "def f(*args, **kwargs):\n pass\n" + nodes, _ = _run(tmp_path, src) + params = _nodes_of_kind(nodes, NodeKind.PARAMETER) + names = {p.name for p in params} + assert "*args" in names + assert "**kwargs" in names + + +# --------------------------------------------------------------------------- +# DATA_ASSIGN: param → variable +# --------------------------------------------------------------------------- + + +def test_param_to_variable_assign(tmp_path: Path) -> None: + """def f(a): x = a → DATA_ASSIGN from PARAMETER a to VARIABLE x.""" + src = "def f(a):\n x = a\n" + nodes, edges = _run(tmp_path, src) + a_param = _param_node(nodes, "a") + x_var = _var_node(nodes, "x") + assert a_param is not None + assert x_var is not None + assigns = _data_edges(edges, EdgeKind.DATA_ASSIGN) + matching = [e for e in assigns if e.src == a_param.id and e.dst == x_var.id] + assert len(matching) == 1 + + +# --------------------------------------------------------------------------- +# DATA_ARG: positional and keyword +# --------------------------------------------------------------------------- + + +def test_data_arg_positional(tmp_path: Path) -> None: + """f calls g(x) → DATA_ARG from VARIABLE x with position 0.""" + src = "def f(a):\n x = a\n g(x)\n" + nodes, edges = _run(tmp_path, src) + x_var = _var_node(nodes, "x") + assert x_var is not None + args = _data_edges(edges, EdgeKind.DATA_ARG) + matching = [ + e for e in args + if e.src == x_var.id and e.metadata.get("position") == 0 + ] + assert len(matching) == 1 + assert matching[0].metadata["callee"] == "g" + + +def test_data_arg_kwarg(tmp_path: Path) -> None: + """g(user=x) → DATA_ARG with kwarg='user'.""" + src = "def f(a):\n x = a\n g(user=x)\n" + nodes, edges = _run(tmp_path, src) + x_var = _var_node(nodes, "x") + assert x_var is not None + args = _data_edges(edges, EdgeKind.DATA_ARG) + matching = [ + e for e in args + if e.src == x_var.id and e.metadata.get("kwarg") == "user" + ] + assert len(matching) == 1 + assert matching[0].dst == "unresolved::arg::g::user" + + +def test_data_arg_param_directly(tmp_path: Path) -> None: + """Passing a parameter directly: def f(a): g(a).""" + src = "def f(a):\n g(a)\n" + nodes, edges = _run(tmp_path, src) + a_param = _param_node(nodes, "a") + assert a_param is not None + args = _data_edges(edges, EdgeKind.DATA_ARG) + matching = [e for e in args if e.src == a_param.id] + assert len(matching) >= 1 + assert matching[0].metadata["position"] == 0 + + +# --------------------------------------------------------------------------- +# DATA_RETURN +# --------------------------------------------------------------------------- + + +def test_data_return_from_variable(tmp_path: Path) -> None: + """return x → DATA_RETURN from VARIABLE x to sentinel.""" + src = "def f(a):\n x = a\n return x\n" + nodes, edges = _run(tmp_path, src) + x_var = _var_node(nodes, "x") + ret_sentinel = _return_sentinel(nodes, "sample.f") + assert x_var is not None + assert ret_sentinel is not None + returns = _data_edges(edges, EdgeKind.DATA_RETURN) + matching = [e for e in returns if e.src == x_var.id and e.dst == ret_sentinel.id] + assert len(matching) == 1 + + +def test_data_return_from_param_directly(tmp_path: Path) -> None: + """return a → DATA_RETURN directly from PARAMETER a.""" + src = "def f(a):\n return a\n" + nodes, edges = _run(tmp_path, src) + a_param = _param_node(nodes, "a") + assert a_param is not None + returns = _data_edges(edges, EdgeKind.DATA_RETURN) + assert any(e.src == a_param.id for e in returns) + + +# --------------------------------------------------------------------------- +# Call-result assignment: y = g(x) → DATA_ASSIGN from unresolved::ret::g +# --------------------------------------------------------------------------- + + +def test_call_result_data_assign(tmp_path: Path) -> None: + """y = g(x) → DATA_ASSIGN from unresolved::ret::g to VARIABLE y.""" + src = "def f(a):\n y = g(a)\n" + nodes, edges = _run(tmp_path, src) + y_var = _var_node(nodes, "y") + assert y_var is not None + assigns = _data_edges(edges, EdgeKind.DATA_ASSIGN) + ret_assigns = [ + e for e in assigns + if e.dst == y_var.id and e.src == "unresolved::ret::g" + ] + assert len(ret_assigns) == 1 + assert ret_assigns[0].metadata["callee"] == "g" + + +# --------------------------------------------------------------------------- +# Shadowing +# --------------------------------------------------------------------------- + + +def test_shadowing_produces_distinct_variable_nodes(tmp_path: Path) -> None: + """x = a; x = b → two distinct VARIABLE nodes (line-qualified).""" + src = "def f(a, b):\n x = a\n x = b\n" + nodes, _edges = _run(tmp_path, src) + x_vars = [ + n for n in nodes + if n.kind == NodeKind.VARIABLE and n.name == "x" + ] + assert len(x_vars) == 2 + # They should have different qualnames (line-qualified). + assert x_vars[0].qualname != x_vars[1].qualname + + +def test_shadowing_edges_point_correctly(tmp_path: Path) -> None: + """Second assignment x = b should use PARAMETER b as source, not first x.""" + src = "def f(a, b):\n x = a\n x = b\n" + nodes, edges = _run(tmp_path, src) + b_param = _param_node(nodes, "b") + a_param = _param_node(nodes, "a") + assert b_param is not None and a_param is not None + + assigns = _data_edges(edges, EdgeKind.DATA_ASSIGN) + # Find DATA_ASSIGN where source is PARAMETER b + b_assigns = [e for e in assigns if e.src == b_param.id] + assert len(b_assigns) >= 1 + # Find DATA_ASSIGN where source is PARAMETER a + a_assigns = [e for e in assigns if e.src == a_param.id] + assert len(a_assigns) >= 1 + # Destinations should be different VARIABLE nodes + assert b_assigns[0].dst != a_assigns[0].dst + + +# --------------------------------------------------------------------------- +# Nested function: inner vars do NOT leak into outer scope +# --------------------------------------------------------------------------- + + +def test_nested_def_vars_do_not_leak(tmp_path: Path) -> None: + """Inner function's variables are not added to the outer function's scope.""" + src = ( + "def outer(a):\n" + " def inner(b):\n" + " z = b\n" + " return a\n" + ) + nodes, edges = _run(tmp_path, src) + # outer should have a return sentinel (from 'return a') + outer_ret = _return_sentinel(nodes, "sample.outer") + assert outer_ret is not None + + # 'z' is inner's variable; outer's DATA_RETURN should only reference 'a' + a_param = _param_node(nodes, "a") + assert a_param is not None + returns = _data_edges(edges, EdgeKind.DATA_RETURN) + outer_returns = [e for e in returns if e.dst == outer_ret.id] + assert len(outer_returns) >= 1 + # All sources of outer's return should be from 'a' (not 'z' or 'b') + for e in outer_returns: + assert e.src == a_param.id + + +def test_inner_function_has_its_own_params(tmp_path: Path) -> None: + """Inner function emits its own PARAMETER nodes.""" + src = ( + "def outer(a):\n" + " def inner(b):\n" + " pass\n" + ) + nodes, _ = _run(tmp_path, src) + params = _nodes_of_kind(nodes, NodeKind.PARAMETER) + param_names = {p.name for p in params} + assert "a" in param_names + assert "b" in param_names + + +# --------------------------------------------------------------------------- +# Kill-switch: emit_data_edges = False +# --------------------------------------------------------------------------- + + +def test_kill_switch_suppresses_data_edges(tmp_path: Path) -> None: + """PythonExtractor.emit_data_edges=False suppresses all C1 edges/nodes.""" + src = "def f(a):\n x = a\n return x\n" + target = tmp_path / "sample.py" + target.write_text(src) + extractor = PythonExtractor() + extractor.__class__.emit_data_edges = False + try: + nodes, edges = extractor.parse_file(target, tmp_path) + finally: + extractor.__class__.emit_data_edges = True # restore + + # No PARAMETER nodes + params = _nodes_of_kind(nodes, NodeKind.PARAMETER) + assert len(params) == 0 + # No data edges + for kind in (EdgeKind.DATA_ASSIGN, EdgeKind.DATA_ARG, EdgeKind.DATA_RETURN, EdgeKind.PARAM_OF): + assert len(_data_edges(edges, kind)) == 0 + + +# --------------------------------------------------------------------------- +# Perf sanity: parse a larger file (codegraph/cli.py) +# --------------------------------------------------------------------------- + + +def test_parse_large_file_perf(tmp_path: Path) -> None: + """Parsing a larger file completes in reasonable time. + + This is a rough sanity check — not a strict performance gate. + The codegraph/cli.py is used as the test fixture (copied to tmp). + """ + # Find cli.py relative to this repo. + repo_root = Path(__file__).parent.parent + cli_py = repo_root / "codegraph" / "cli.py" + if not cli_py.exists(): + pytest.skip("cli.py not found — skipping perf sanity") + + dest = tmp_path / "cli.py" + shutil.copy(cli_py, dest) + + extractor = PythonExtractor() + + # Baseline: without data edges + extractor.__class__.emit_data_edges = False + t0 = time.perf_counter() + extractor.parse_file(dest, tmp_path) + t_without = time.perf_counter() - t0 + + # With data edges + extractor.__class__.emit_data_edges = True + t1 = time.perf_counter() + nodes_with, edges_with = extractor.parse_file(dest, tmp_path) + t_with = time.perf_counter() - t1 + + data_edge_count = sum( + 1 for e in edges_with + if e.kind in (EdgeKind.DATA_ASSIGN, EdgeKind.DATA_ARG, EdgeKind.DATA_RETURN) + ) + param_node_count = sum(1 for n in nodes_with if n.kind == NodeKind.PARAMETER) + + # Must stay in the same order of magnitude (not 10x slower) + assert t_with < max(t_without * 10, 1.0), ( + f"Data edge collection too slow: {t_without:.3f}s without vs {t_with:.3f}s with" + ) + + # Sanity: should produce data edges on a real file + assert data_edge_count > 0, "Expected data edges on cli.py" + assert param_node_count > 0, "Expected PARAMETER nodes on cli.py" + + +def test_for_loop_target_fed_by_iterable(tmp_path: Path) -> None: + """``for url in scraped_urls`` → DATA_ASSIGN from iterable var to loop var.""" + nodes, edges = _run( + tmp_path, + "def f(scraped_urls):\n" + " for url in scraped_urls:\n" + " fetch(url)\n", + ) + loop_vars = [ + n for n in nodes + if n.kind == NodeKind.VARIABLE and n.name == "url" + ] + assert len(loop_vars) == 1 + assert loop_vars[0].metadata.get("loop_target") is True + # iterable (PARAMETER scraped_urls) feeds the loop var + param = next( + n for n in nodes + if n.kind == NodeKind.PARAMETER and n.name == "scraped_urls" + ) + assign = [ + e for e in edges + if e.kind == EdgeKind.DATA_ASSIGN + and e.src == param.id and e.dst == loop_vars[0].id + ] + assert assign, "expected DATA_ASSIGN from iterable param to loop target" + # and the loop var flows into the fetch() call argument + arg = [ + e for e in edges + if e.kind == EdgeKind.DATA_ARG and e.src == loop_vars[0].id + ] + assert arg, "expected DATA_ARG from loop var into fetch()"