fix(python): restore method call edges in the call graph

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<String, NodeId> = 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<String, Vec<String>> = 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<String> = 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:?}"
+        );
+    }
 }

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::<Vec<_>>()
+        );
+        for class in &ir.classes {
+            println!(
+                "CLASS {:?} methods: {:?}",
+                class.name,
+                class.methods.iter().map(|m| &m.name).collect::<Vec<_>>()
+            );
+        }
+
+        // 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#"