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:?}" + ); + } } 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() {