Skip to main content

mollify_core/
installed.rs

1//! Installed-environment introspection. When a virtualenv is present, reads
2//! `*.dist-info` metadata from `site-packages` to (a) map import names to
3//! distributions accurately (beyond the static alias table) and (b) know which
4//! distributions are actually installed — which lets `deps` distinguish a
5//! **transitive** dependency (installed but undeclared) from a genuinely
6//! **missing** one (not installed at all). Best-effort: absent venv → `None`.
7
8use crate::known::normalize_dist;
9use camino::Utf8Path;
10use rustc_hash::{FxHashMap, FxHashSet};
11
12#[derive(Debug, Default)]
13pub struct Installed {
14    /// import top-level name → normalized distribution name.
15    pub import_to_dist: FxHashMap<String, String>,
16    /// All installed (normalized) distribution names.
17    pub dists: FxHashSet<String>,
18    /// normalized distribution name → installed version (from dist-info METADATA).
19    /// Lets supply-chain resolve a declared *range* to the concrete version that
20    /// is actually installed, for precise advisory matching.
21    pub versions: FxHashMap<String, String>,
22}
23
24/// Discover and parse the project's virtualenv `site-packages`, if any.
25pub fn discover(root: &Utf8Path) -> Option<Installed> {
26    let sp = find_site_packages(root)?;
27    let mut inst = Installed::default();
28    for entry in std::fs::read_dir(&sp).ok()?.flatten() {
29        let name = entry.file_name();
30        let name = name.to_string_lossy();
31        if !name.ends_with(".dist-info") {
32            continue;
33        }
34        let dir = entry.path();
35        // Distribution name from METADATA `Name:`, else the dir prefix.
36        let meta = std::fs::read_to_string(dir.join("METADATA")).ok();
37        let dist = meta
38            .as_ref()
39            .and_then(|m| {
40                m.lines()
41                    .find_map(|l| l.strip_prefix("Name:").map(|n| n.trim().to_string()))
42            })
43            .unwrap_or_else(|| name.split('-').next().unwrap_or(&name).to_string());
44        let dist = normalize_dist(&dist);
45        inst.dists.insert(dist.clone());
46        if let Some(ver) = meta.as_ref().and_then(|m| {
47            m.lines()
48                .find_map(|l| l.strip_prefix("Version:").map(|v| v.trim().to_string()))
49        }) {
50            inst.versions.insert(dist.clone(), ver);
51        }
52
53        // Import names from top_level.txt; fall back to the dist name.
54        let tops = std::fs::read_to_string(dir.join("top_level.txt"))
55            .ok()
56            .map(|t| {
57                t.lines()
58                    .map(|l| l.trim().to_string())
59                    .filter(|l| !l.is_empty())
60                    .collect::<Vec<_>>()
61            })
62            .filter(|v| !v.is_empty())
63            .unwrap_or_else(|| vec![dist.replace('-', "_")]);
64        for top in tops {
65            // Only the package's top segment matters for import resolution.
66            let top = top.split('/').next().unwrap_or(&top).to_string();
67            inst.import_to_dist
68                .entry(top)
69                .or_insert_with(|| dist.clone());
70        }
71    }
72    if inst.dists.is_empty() {
73        None
74    } else {
75        Some(inst)
76    }
77}
78
79/// Locate a `site-packages` directory for the project (common venv layouts and
80/// `$VIRTUAL_ENV`). Returns the first that exists.
81fn find_site_packages(root: &Utf8Path) -> Option<camino::Utf8PathBuf> {
82    let mut roots: Vec<camino::Utf8PathBuf> = vec![
83        root.join(".venv"),
84        root.join("venv"),
85        root.join("env"),
86        root.join(".env"),
87    ];
88    if let Ok(v) = std::env::var("VIRTUAL_ENV") {
89        roots.insert(0, camino::Utf8PathBuf::from(v));
90    }
91    for venv in roots {
92        // POSIX: <venv>/lib/pythonX.Y/site-packages ; Windows: <venv>/Lib/site-packages
93        for libdir in ["lib", "Lib"] {
94            let base = venv.join(libdir);
95            let Ok(rd) = std::fs::read_dir(&base) else {
96                continue;
97            };
98            // Windows layout: Lib/site-packages directly.
99            let direct = base.join("site-packages");
100            if direct.is_dir() {
101                return Some(direct);
102            }
103            // POSIX: lib/pythonX.Y/site-packages.
104            for e in rd.flatten() {
105                let p = e.path();
106                if p.is_dir() {
107                    if let Ok(p) = camino::Utf8PathBuf::from_path_buf(p) {
108                        let sp = p.join("site-packages");
109                        if sp.is_dir() {
110                            return Some(sp);
111                        }
112                    }
113                }
114            }
115        }
116    }
117    None
118}
119
120#[cfg(test)]
121mod tests {
122    use super::*;
123    use camino::Utf8PathBuf;
124
125    #[test]
126    fn parses_dist_info_from_a_synthetic_venv() {
127        let d = std::env::temp_dir().join(format!("mollify-installed-{}", std::process::id()));
128        let _ = std::fs::remove_dir_all(&d);
129        let sp = d.join(".venv/lib/python3.11/site-packages/requests-2.31.0.dist-info");
130        std::fs::create_dir_all(&sp).unwrap();
131        std::fs::write(sp.join("METADATA"), "Name: requests\nVersion: 2.31.0\n").unwrap();
132        std::fs::write(sp.join("top_level.txt"), "requests\n").unwrap();
133        let root = Utf8PathBuf::from_path_buf(d.clone()).unwrap();
134        let inst = discover(&root).unwrap();
135        assert!(inst.dists.contains("requests"));
136        assert_eq!(
137            inst.import_to_dist.get("requests").map(|s| s.as_str()),
138            Some("requests")
139        );
140        std::fs::remove_dir_all(&d).ok();
141    }
142}