From 9e586f65174315dfd5676fd6aeac3c493a2c812c Mon Sep 17 00:00:00 2001 From: zhanjisheng Date: Fri, 23 Oct 2026 23:22:22 +0800 Subject: [PATCH 1/2] fix(python): restore method call edges in the call graph The Python parser dropped every method call (`self.x()`, `obj.method()`), so `get_callers`, `get_callees`, `analyze_impact` and `find_dead_imports` were blind on Python codebases. Two defects had to be fixed together: 1. `extract_callee_name` returned the full text of a tree-sitter `attribute` node (`"self.simulator.partial_close"`). No symbol is registered under that name, so every callee lookup missed. It now keeps only the trailing method name, mirroring how `codegraph-rust` handles `field_expression`. 2. `ir_to_graph` indexed its `node_map` by bare name only, while the extractor attributes calls made inside a method to the qualified name `"Class.method"`. Every method call therefore missed its caller lookup and was silently dropped. The qualified name is now registered alongside the bare one. Before: `run_tick` in a Python trading codebase reported 0 callees and `analyze_impact` rated a live-trading hot path as `risk: low`. After: 57 callees, 47 callers, `risk: medium`, and no false dead-import reports. Regression test added over the same path the server uses (`parse_source` -> `ir_to_graph`). --- crates/codegraph-python/src/extractor.rs | 69 +++++++++++++++++++- crates/codegraph-python/src/parser_impl.rs | 74 ++++++++++++++++++++++ 2 files changed, 142 insertions(+), 1 deletion(-) diff --git a/crates/codegraph-python/src/extractor.rs b/crates/codegraph-python/src/extractor.rs index 4733abc..a56bdbc 100644 --- a/crates/codegraph-python/src/extractor.rs +++ b/crates/codegraph-python/src/extractor.rs @@ -610,7 +610,18 @@ fn extract_callee_name(source: &[u8], node: Node) -> String { "identifier" => node.utf8_text(source).unwrap_or("").to_string(), "attribute" => { // Handle obj.method() or self.method() - node.utf8_text(source).unwrap_or("").to_string() + // + // Keep only the trailing method name. The receiver (`self`, + // `self.simulator`, `np`, ...) is not part of the callee's identity, + // and keeping it produced callee keys like "self.simulator.partial_close" + // that matched no symbol in the graph — so every method call silently + // lost its edge. Mirrors codegraph-rust's `field_expression` handling, + // which keeps just the field name. + let full = node.utf8_text(source).unwrap_or(""); + match full.rsplit_once('.') { + Some((_, method)) => method.to_string(), + None => full.to_string(), + } } _ => String::new(), } @@ -941,6 +952,62 @@ def greet(name): assert!(ir.calls.len() >= 2, "Should find at least 2 calls"); } + #[test] + fn test_method_call_callee_names() { + // Method calls must be recorded with the bare method name so they can + // be matched against the symbol map; the receiver (`self`, + // `self.simulator`) is not part of the callee's identity. + let source = r#" +class TradingEngine: + def run_tick(self): + self.simulator.close_position("BTC", 1.0) + self._emit_ops("x") + self._pool_ranks(alt_dfs) + + def _emit_ops(self, x): + pass + + def _pool_ranks(self, alt_dfs): + pass +"#; + let path = Path::new("test.py"); + let config = ParserConfig::default(); + let ir = extract(source, path, &config).unwrap(); + + let calls: Vec<(&str, &str)> = ir + .calls + .iter() + .map(|c| (c.caller.as_str(), c.callee.as_str())) + .collect(); + println!("extracted calls: {calls:?}"); + + println!( + "module-level functions: {:?}", + ir.functions.iter().map(|f| &f.name).collect::>() + ); + for class in &ir.classes { + println!( + "CLASS {:?} methods: {:?}", + class.name, + class.methods.iter().map(|m| &m.name).collect::>() + ); + } + + // Every call inside a method is attributed to the qualified method name. + assert!( + calls.iter().all(|(caller, _)| *caller == "TradingEngine.run_tick"), + "all calls should be attributed to TradingEngine.run_tick, got {calls:?}" + ); + + // The receiver is stripped; only the method name survives. + for expected in ["close_position", "_emit_ops", "_pool_ranks"] { + assert!( + calls.iter().any(|(_, callee)| *callee == expected), + "expected a call to {expected}, got {calls:?}" + ); + } + } + #[test] fn test_extract_imports() { let source = r#" diff --git a/crates/codegraph-python/src/parser_impl.rs b/crates/codegraph-python/src/parser_impl.rs index 7980978..6ec2631 100644 --- a/crates/codegraph-python/src/parser_impl.rs +++ b/crates/codegraph-python/src/parser_impl.rs @@ -201,6 +201,12 @@ impl PythonParser { graph .add_edge(class_id, method_id, EdgeType::Contains, PropertyMap::new()) .map_err(|e| ParserError::GraphError(e.to_string()))?; + // Also index the method under its qualified name. The + // extractor attributes calls made inside a method to + // "Class.method", while node_map is keyed by bare names — + // without this entry every method call failed to resolve + // and its edge was silently dropped. + node_map.insert(format!("{}.{}", class.name, method.name), method_id); } } } @@ -554,6 +560,74 @@ fn extract_first_string_arg(attr: &str) -> Option { #[cfg(test)] mod tests { use super::*; + use codegraph::{EdgeType, NodeType}; + + #[test] + fn test_method_call_edge_is_created() { + // End-to-end over the same path the server uses (parse_source -> + // ir_to_graph). A call made inside a method must produce a Calls edge. + let source = r#" +class Probe: + def probe_method(self): + return 1 + + def probe_method_caller(self): + return self.probe_method() + + +def probe_free_fn(): + return 1 + + +def probe_free_caller(): + return probe_free_fn() +"#; + let mut graph = CodeGraph::in_memory().unwrap(); + let parser = PythonParser::new(); + parser + .parse_source(source, std::path::Path::new("probe.py"), &mut graph) + .expect("parse source"); + + let funcs = graph + .query() + .node_type(NodeType::Function) + .execute() + .expect("query functions"); + + let find = |name: &str| { + funcs + .iter() + .find(|id| { + graph + .get_node(**id) + .ok() + .and_then(|n| n.properties.get_string("name").map(|n| n == name)) + .unwrap_or(false) + }) + .copied() + .expect("node for function") + }; + + for (caller_name, callee_name) in [ + ("probe_method_caller", "probe_method"), + ("probe_free_caller", "probe_free_fn"), + ] { + let caller = find(caller_name); + let callee = find(callee_name); + let edges = graph + .get_edges_between(caller, callee) + .expect("edges between caller and callee"); + assert!( + edges.iter().any(|e| { + graph + .get_edge(*e) + .map(|edge| edge.edge_type == EdgeType::Calls) + .unwrap_or(false) + }), + "expected a Calls edge from {caller_name} to {callee_name}" + ); + } + } #[test] fn test_python_parser_new() { From 64a9525338e431a40349d059ec09e02bfe1dc120 Mon Sep 17 00:00:00 2001 From: zhanjisheng Date: Fri, 23 Oct 2026 23:22:34 +0800 Subject: [PATCH 2/2] fix(python): apply the same call-graph fix to the parallel builder MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit `builder::build_graph` is a second implementation of the same job as `parser_impl::ir_to_graph` (the former is reached only through `parser.rs`'s `parse_source`). It carried the same two defects as the server path, so it is brought to parity here. It also had a defect of its own: the extractor lists a class's methods in both `ir.functions` and `ir.classes[].methods`, so every method was built twice — once as a stray module-level node and once as a real method node linked to its class. The duplicate left two nodes with the same name and split the call edges away from the node that name lookups found. Module function entries that are really class methods are now skipped (matched on name and line span, so a same-named module function is kept). The bare method name is registered alongside the qualified one so a call like `self.helper()` resolves within its own file instead of relying on the server's global name lookup, which cannot distinguish two classes that share a method name. --- crates/codegraph-python/src/builder.rs | 224 ++++++++++++++++++++++++- 1 file changed, 219 insertions(+), 5 deletions(-) diff --git a/crates/codegraph-python/src/builder.rs b/crates/codegraph-python/src/builder.rs index 2ae950c..81b9f30 100644 --- a/crates/codegraph-python/src/builder.rs +++ b/crates/codegraph-python/src/builder.rs @@ -28,8 +28,37 @@ pub fn build_graph(graph: &mut CodeGraph, ir: &CodeIR, file_path: &str) -> Resul // Track entity name -> NodeId mappings for relationship building let mut entity_map: HashMap = HashMap::new(); + // The extractor lists a class's methods in BOTH `ir.functions` and + // `ir.classes[].methods` (see `apply_class_extraction`). Without this + // filter every method was built twice — once as a stray module-level + // function node and once as a real method node linked to its class. The + // duplicate left the graph with two nodes carrying the same name, and the + // call edges landed on the method node while lookups by name hit the + // orphan, so callers/callees looked empty. Match on name AND line span so + // a module-level function that merely shares a name with a method is kept. + let method_spans: std::collections::HashSet<(&str, i64, i64)> = ir + .classes + .iter() + .flat_map(|class| class.methods.iter()) + .map(|method| { + ( + method.name.as_str(), + method.line_start as i64, + method.line_end as i64, + ) + }) + .collect(); + // Add all functions for func in &ir.functions { + // This function is a class method — the class pass below owns it. + if method_spans.contains(&( + func.name.as_str(), + func.line_start as i64, + func.line_end as i64, + )) { + continue; + } let mut props = PropertyMap::new() .with("name", func.name.clone()) .with("signature", func.signature.clone()) @@ -110,16 +139,53 @@ pub fn build_graph(graph: &mut CodeGraph, ir: &CodeIR, file_path: &str) -> Resul // Track methods with qualified name for call relationships let qualified = format!("{}.{}", class.name, method.name); entity_map.insert(qualified, method_id); + // Also register the bare method name so a call like + // `self.helper()` inside the same file resolves directly, instead + // of relying on the server's global name lookup (which cannot + // distinguish two classes that share a method name). + entity_map.entry(method.name.clone()).or_insert(method_id); } } // Add call relationships + // + // A callee that lives in another file is not in this file's entity map, so + // its edge cannot be drawn here. Those calls are parked on the caller node's + // `unresolved_calls` property and connected later by the server's + // cross-file resolution pass (GraphUpdater::resolve_cross_file_imports). + // Same pattern as codegraph-fortran / codegraph-rust / codegraph-swift. + let mut unresolved_calls: HashMap> = HashMap::new(); + for call in &ir.calls { - if let (Some(&caller_id), Some(&callee_id)) = - (entity_map.get(&call.caller), entity_map.get(&call.callee)) - { - helpers::add_call(graph, caller_id, callee_id, call.call_site_line as i64) - .map_err(|e| crate::error::ParseError::GraphError(e.to_string()))?; + if let Some(&caller_id) = entity_map.get(&call.caller) { + if let Some(&callee_id) = entity_map.get(&call.callee) { + helpers::add_call(graph, caller_id, callee_id, call.call_site_line as i64) + .map_err(|e| crate::error::ParseError::GraphError(e.to_string()))?; + } else { + unresolved_calls + .entry(call.caller.clone()) + .or_default() + .push(call.callee.clone()); + } + } + } + + // Store unresolved calls on caller nodes for cross-file resolution + for (caller_name, callees) in unresolved_calls { + if let Some(&caller_id) = entity_map.get(&caller_name) { + if let Ok(node) = graph.get_node(caller_id) { + let mut all_callees: Vec = node + .properties + .get_string_list_compat("unresolved_calls") + .unwrap_or_default(); + for callee in &callees { + if !all_callees.iter().any(|c| c == callee) { + all_callees.push(callee.clone()); + } + } + let new_props = node.properties.clone().with("unresolved_calls", all_callees); + let _ = graph.update_node_properties(caller_id, new_props); + } } } @@ -214,4 +280,152 @@ mod tests { let result = build_graph(&mut graph, &ir, "test.py"); assert!(result.is_ok()); } + + #[test] + fn test_local_call_creates_edge() { + let mut graph = CodeGraph::in_memory().unwrap(); + let mut ir = CodeIR::new(std::path::PathBuf::from("test.py")); + + ir.add_function(FunctionEntity::new("caller", 1, 3)); + ir.add_function(FunctionEntity::new("callee", 5, 7)); + ir.add_call(CallRelation::new("caller", "callee", 2)); + + build_graph(&mut graph, &ir, "test.py").unwrap(); + + let funcs = graph + .query() + .node_type(NodeType::Function) + .execute() + .expect("query functions"); + assert_eq!(funcs.len(), 2); + + let caller = funcs + .iter() + .find(|id| { + graph + .get_node(**id) + .ok() + .and_then(|n| n.properties.get_string("name").map(|n| n == "caller")) + .unwrap_or(false) + }) + .copied() + .expect("caller node"); + let callee = funcs + .iter() + .find(|id| { + graph + .get_node(**id) + .ok() + .and_then(|n| n.properties.get_string("name").map(|n| n == "callee")) + .unwrap_or(false) + }) + .copied() + .expect("callee node"); + + let edges = graph + .get_edges_between(caller, callee) + .expect("edges between caller and callee"); + assert!( + edges.iter().any(|e| { + graph + .get_edge(*e) + .map(|edge| edge.edge_type == EdgeType::Calls) + .unwrap_or(false) + }), + "expected a Calls edge between caller and callee" + ); + } + + #[test] + fn test_class_method_yields_single_node_and_call_edge() { + // The extractor lists a class's methods in both `ir.functions` and + // `ir.classes[].methods`. The builder must not build a stray + // module-level copy of every method — two nodes with the same name + // split the call edges away from the node that lookups find. + let source = r#" +class TradingEngine: + def run_tick(self): + self._emit_ops("x") + + def _emit_ops(self, x): + pass +"#; + let config = crate::config::ParserConfig::default(); + let ir = crate::extractor::extract(source, std::path::Path::new("test.py"), &config) + .expect("extract source"); + + let mut graph = CodeGraph::in_memory().unwrap(); + build_graph(&mut graph, &ir, "test.py").unwrap(); + + let funcs = graph + .query() + .node_type(NodeType::Function) + .execute() + .expect("query functions"); + assert_eq!( + funcs.len(), + 2, + "each method must yield exactly one node, got {}", + funcs.len() + ); + + let find = |name: &str| { + funcs + .iter() + .find(|id| { + graph + .get_node(**id) + .ok() + .and_then(|n| n.properties.get_string("name").map(|n| n == name)) + .unwrap_or(false) + }) + .copied() + .expect("node for method") + }; + let caller = find("run_tick"); + let callee = find("_emit_ops"); + + let edges = graph + .get_edges_between(caller, callee) + .expect("edges between run_tick and _emit_ops"); + assert!( + edges.iter().any(|e| { + graph + .get_edge(*e) + .map(|edge| edge.edge_type == EdgeType::Calls) + .unwrap_or(false) + }), + "expected a Calls edge from run_tick to _emit_ops" + ); + } + + #[test] + fn test_unresolved_call_is_parked_for_cross_file_resolution() { let mut graph = CodeGraph::in_memory().unwrap(); + let mut ir = CodeIR::new(std::path::PathBuf::from("test.py")); + + ir.add_function(FunctionEntity::new("caller", 1, 3)); + // The callee lives in another file, so it is not in this file's entity + // map. It must be parked on the caller node as `unresolved_calls` so the + // server's cross-file pass can connect it later. + ir.add_call(CallRelation::new("caller", "partial_close", 2)); + + build_graph(&mut graph, &ir, "test.py").unwrap(); + + let funcs = graph + .query() + .node_type(NodeType::Function) + .execute() + .expect("query functions"); + assert_eq!(funcs.len(), 1); + + let node = graph.get_node(funcs[0]).expect("caller node"); + let unresolved = node + .properties + .get_string_list_compat("unresolved_calls") + .unwrap_or_default(); + assert!( + unresolved.iter().any(|c| c == "partial_close"), + "callee should be parked for cross-file resolution, got {unresolved:?}" + ); + } }