Skip to content

Commit 9e586f6

Browse files
committed
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`).
1 parent 590c2ca commit 9e586f6

2 files changed

Lines changed: 142 additions & 1 deletion

File tree

‎crates/codegraph-python/src/extractor.rs‎

Lines changed: 68 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -610,7 +610,18 @@ fn extract_callee_name(source: &[u8], node: Node) -> String {
610610
"identifier" => node.utf8_text(source).unwrap_or("").to_string(),
611611
"attribute" => {
612612
// Handle obj.method() or self.method()
613-
node.utf8_text(source).unwrap_or("").to_string()
613+
//
614+
// Keep only the trailing method name. The receiver (`self`,
615+
// `self.simulator`, `np`, ...) is not part of the callee's identity,
616+
// and keeping it produced callee keys like "self.simulator.partial_close"
617+
// that matched no symbol in the graph — so every method call silently
618+
// lost its edge. Mirrors codegraph-rust's `field_expression` handling,
619+
// which keeps just the field name.
620+
let full = node.utf8_text(source).unwrap_or("");
621+
match full.rsplit_once('.') {
622+
Some((_, method)) => method.to_string(),
623+
None => full.to_string(),
624+
}
614625
}
615626
_ => String::new(),
616627
}
@@ -941,6 +952,62 @@ def greet(name):
941952
assert!(ir.calls.len() >= 2, "Should find at least 2 calls");
942953
}
943954

955+
#[test]
956+
fn test_method_call_callee_names() {
957+
// Method calls must be recorded with the bare method name so they can
958+
// be matched against the symbol map; the receiver (`self`,
959+
// `self.simulator`) is not part of the callee's identity.
960+
let source = r#"
961+
class TradingEngine:
962+
def run_tick(self):
963+
self.simulator.close_position("BTC", 1.0)
964+
self._emit_ops("x")
965+
self._pool_ranks(alt_dfs)
966+
967+
def _emit_ops(self, x):
968+
pass
969+
970+
def _pool_ranks(self, alt_dfs):
971+
pass
972+
"#;
973+
let path = Path::new("test.py");
974+
let config = ParserConfig::default();
975+
let ir = extract(source, path, &config).unwrap();
976+
977+
let calls: Vec<(&str, &str)> = ir
978+
.calls
979+
.iter()
980+
.map(|c| (c.caller.as_str(), c.callee.as_str()))
981+
.collect();
982+
println!("extracted calls: {calls:?}");
983+
984+
println!(
985+
"module-level functions: {:?}",
986+
ir.functions.iter().map(|f| &f.name).collect::<Vec<_>>()
987+
);
988+
for class in &ir.classes {
989+
println!(
990+
"CLASS {:?} methods: {:?}",
991+
class.name,
992+
class.methods.iter().map(|m| &m.name).collect::<Vec<_>>()
993+
);
994+
}
995+
996+
// Every call inside a method is attributed to the qualified method name.
997+
assert!(
998+
calls.iter().all(|(caller, _)| *caller == "TradingEngine.run_tick"),
999+
"all calls should be attributed to TradingEngine.run_tick, got {calls:?}"
1000+
);
1001+
1002+
// The receiver is stripped; only the method name survives.
1003+
for expected in ["close_position", "_emit_ops", "_pool_ranks"] {
1004+
assert!(
1005+
calls.iter().any(|(_, callee)| *callee == expected),
1006+
"expected a call to {expected}, got {calls:?}"
1007+
);
1008+
}
1009+
}
1010+
9441011
#[test]
9451012
fn test_extract_imports() {
9461013
let source = r#"

‎crates/codegraph-python/src/parser_impl.rs‎

Lines changed: 74 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -201,6 +201,12 @@ impl PythonParser {
201201
graph
202202
.add_edge(class_id, method_id, EdgeType::Contains, PropertyMap::new())
203203
.map_err(|e| ParserError::GraphError(e.to_string()))?;
204+
// Also index the method under its qualified name. The
205+
// extractor attributes calls made inside a method to
206+
// "Class.method", while node_map is keyed by bare names —
207+
// without this entry every method call failed to resolve
208+
// and its edge was silently dropped.
209+
node_map.insert(format!("{}.{}", class.name, method.name), method_id);
204210
}
205211
}
206212
}
@@ -554,6 +560,74 @@ fn extract_first_string_arg(attr: &str) -> Option<String> {
554560
#[cfg(test)]
555561
mod tests {
556562
use super::*;
563+
use codegraph::{EdgeType, NodeType};
564+
565+
#[test]
566+
fn test_method_call_edge_is_created() {
567+
// End-to-end over the same path the server uses (parse_source ->
568+
// ir_to_graph). A call made inside a method must produce a Calls edge.
569+
let source = r#"
570+
class Probe:
571+
def probe_method(self):
572+
return 1
573+
574+
def probe_method_caller(self):
575+
return self.probe_method()
576+
577+
578+
def probe_free_fn():
579+
return 1
580+
581+
582+
def probe_free_caller():
583+
return probe_free_fn()
584+
"#;
585+
let mut graph = CodeGraph::in_memory().unwrap();
586+
let parser = PythonParser::new();
587+
parser
588+
.parse_source(source, std::path::Path::new("probe.py"), &mut graph)
589+
.expect("parse source");
590+
591+
let funcs = graph
592+
.query()
593+
.node_type(NodeType::Function)
594+
.execute()
595+
.expect("query functions");
596+
597+
let find = |name: &str| {
598+
funcs
599+
.iter()
600+
.find(|id| {
601+
graph
602+
.get_node(**id)
603+
.ok()
604+
.and_then(|n| n.properties.get_string("name").map(|n| n == name))
605+
.unwrap_or(false)
606+
})
607+
.copied()
608+
.expect("node for function")
609+
};
610+
611+
for (caller_name, callee_name) in [
612+
("probe_method_caller", "probe_method"),
613+
("probe_free_caller", "probe_free_fn"),
614+
] {
615+
let caller = find(caller_name);
616+
let callee = find(callee_name);
617+
let edges = graph
618+
.get_edges_between(caller, callee)
619+
.expect("edges between caller and callee");
620+
assert!(
621+
edges.iter().any(|e| {
622+
graph
623+
.get_edge(*e)
624+
.map(|edge| edge.edge_type == EdgeType::Calls)
625+
.unwrap_or(false)
626+
}),
627+
"expected a Calls edge from {caller_name} to {callee_name}"
628+
);
629+
}
630+
}
557631

558632
#[test]
559633
fn test_python_parser_new() {

0 commit comments

Comments
 (0)