Skip to main content

mollify_core/
trace.rs

1//! `mollify trace <module>` — the static dependency neighborhood of a module:
2//! what it imports (callees, "down") and what imports it (callers, "up"). A
3//! lightweight, deterministic answer to "what breaks if I touch this?" built
4//! straight from the import graph (fallow's `trace`, in Python terms).
5
6use mollify_graph::ModuleGraph;
7
8/// The import neighborhood of a target module, both directions, sorted.
9#[derive(Debug, Clone, PartialEq, Eq)]
10pub struct Trace {
11    /// The matched dotted module name (may differ from the user's query).
12    pub target: String,
13    /// Modules the target imports directly.
14    pub imports: Vec<String>,
15    /// Modules that directly import the target.
16    pub imported_by: Vec<String>,
17}
18
19/// Resolve `query` to a module (exact dotted match, else suffix match) and
20/// compute its import neighborhood. `None` if nothing matches.
21pub 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
44/// Exact dotted match first; otherwise the lexicographically-first module whose
45/// dotted name equals the query's trailing segment(s).
46fn 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}