mollify_core/
installed.rs1use crate::known::normalize_dist;
9use camino::Utf8Path;
10use rustc_hash::{FxHashMap, FxHashSet};
11
12#[derive(Debug, Default)]
13pub struct Installed {
14 pub import_to_dist: FxHashMap<String, String>,
16 pub dists: FxHashSet<String>,
18 pub versions: FxHashMap<String, String>,
22}
23
24pub 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 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 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 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
79fn 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 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 let direct = base.join("site-packages");
100 if direct.is_dir() {
101 return Some(direct);
102 }
103 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}