1use mollify_graph::ModuleGraph;
7
8#[derive(Debug, Clone, PartialEq, Eq)]
10pub struct Trace {
11 pub target: String,
13 pub imports: Vec<String>,
15 pub imported_by: Vec<String>,
17}
18
19pub fn module(graph: &ModuleGraph, query: &str) -> Option<Trace> {
22 let target = resolve(graph, query)?;
23 let mut imports = Vec::new();
24 let mut imported_by = Vec::new();
25 for (importer, imported) in graph.import_edges() {
26 if importer == target {
27 imports.push(imported.to_string());
28 }
29 if imported == target {
30 imported_by.push(importer.to_string());
31 }
32 }
33 imports.sort();
34 imports.dedup();
35 imported_by.sort();
36 imported_by.dedup();
37 Some(Trace {
38 target,
39 imports,
40 imported_by,
41 })
42}
43
44fn resolve(graph: &ModuleGraph, query: &str) -> Option<String> {
47 let mut names: Vec<&str> = graph.modules.iter().map(|m| m.dotted.as_str()).collect();
48 names.sort();
49 if let Some(exact) = names.iter().find(|n| **n == query) {
50 return Some((*exact).to_string());
51 }
52 names
53 .iter()
54 .find(|n| n.ends_with(&format!(".{query}")) || **n == query)
55 .map(|n| (*n).to_string())
56}
57
58#[cfg(test)]
59mod tests {
60 use super::*;
61 use camino::Utf8PathBuf;
62 use mollify_graph::discover_python_files;
63
64 fn temp(tag: &str) -> Utf8PathBuf {
65 let base =
66 std::env::temp_dir().join(format!("mollify-core-trace-{}-{tag}", std::process::id()));
67 let _ = std::fs::remove_dir_all(&base);
68 std::fs::create_dir_all(&base).unwrap();
69 Utf8PathBuf::from_path_buf(base).unwrap()
70 }
71
72 #[test]
73 fn traces_both_directions() {
74 let d = temp("t");
75 std::fs::write(d.join("a.py"), "import b\n").unwrap();
76 std::fs::write(d.join("b.py"), "import c\n").unwrap();
77 std::fs::write(d.join("c.py"), "").unwrap();
78 let files = discover_python_files(&d);
79 let g = ModuleGraph::build(&d, &files);
80 let t = module(&g, "b").unwrap();
81 assert_eq!(t.target, "b");
82 assert!(t.imports.contains(&"c".to_string()), "{t:?}");
83 assert!(t.imported_by.contains(&"a".to_string()), "{t:?}");
84 assert!(module(&g, "nonexistent").is_none());
85 std::fs::remove_dir_all(&d).ok();
86 }
87}