1use super::{BloatDir, EnforcePolicy, PackageManager, dir_size, run_command_with_timeout};
14use anyhow::{Result, anyhow};
15use std::collections::{HashMap, HashSet};
16use std::fs;
17use std::path::{Path, PathBuf};
18
19const PYVENV_CFG: &str = "pyvenv.cfg";
21
22const FOREIGN_PYTHON_LOCKFILES: [&str; 3] = ["poetry.lock", "Pipfile.lock", "pdm.lock"];
29
30pub(super) const BASELINE_DISTRIBUTIONS: [&str; 4] =
33 ["pip", "setuptools", "wheel", "pkg-resources"];
34
35pub struct Venv;
37
38fn find_venv_dirs(path: &Path) -> Vec<std::path::PathBuf> {
42 let mut found = Vec::new();
43
44 let Ok(entries) = fs::read_dir(path) else {
45 return found;
46 };
47
48 for entry in entries.flatten() {
49 let entry_path = entry.path();
50 if entry_path.is_dir() && entry_path.join(PYVENV_CFG).exists() {
51 found.push(entry_path);
52 }
53 }
54
55 found
56}
57
58fn is_poetry_project(path: &Path) -> bool {
63 fs::read_to_string(path.join("pyproject.toml"))
64 .map(|c| {
65 c.lines()
66 .any(|l| l.trim_start().starts_with("[tool.poetry"))
67 })
68 .unwrap_or(false)
69}
70
71pub(super) fn normalize_package_name(name: &str) -> String {
74 let mut out = String::with_capacity(name.len());
75 let mut last_dash = false;
76 for c in name.chars() {
77 if c == '-' || c == '_' || c == '.' {
78 if !last_dash {
79 out.push('-');
80 }
81 last_dash = true;
82 } else {
83 out.push(c.to_ascii_lowercase());
84 last_dash = false;
85 }
86 }
87 out
88}
89
90fn requirement_name(spec: &str) -> Option<String> {
96 let name: String = spec
97 .chars()
98 .take_while(|c| c.is_ascii_alphanumeric() || matches!(c, '-' | '_' | '.'))
99 .collect();
100 if name.is_empty() || name.starts_with('.') || spec.contains("://") && !spec.contains(" @ ") {
102 return None;
103 }
104 Some(normalize_package_name(&name))
105}
106
107fn requirement_names(file: &Path, visited: &mut Vec<PathBuf>) -> Option<HashSet<String>> {
113 if visited.len() >= 8 || visited.iter().any(|p| p == file) {
116 return None;
117 }
118 visited.push(file.to_path_buf());
119
120 let content = fs::read_to_string(file).ok()?;
121 let dir = file.parent()?;
122 let mut names = HashSet::new();
123
124 for raw in content.lines() {
125 let mut line = raw.trim();
128 if let Some(idx) = line.find(" #") {
129 line = &line[..idx];
130 }
131 let line = line.trim_end_matches('\\').trim();
132 if line.is_empty() || line.starts_with('#') {
133 continue;
134 }
135
136 if let Some(included) = line
137 .strip_prefix("-r ")
138 .or_else(|| line.strip_prefix("--requirement "))
139 .or_else(|| line.strip_prefix("-c "))
140 .or_else(|| line.strip_prefix("--constraint "))
141 {
142 names.extend(requirement_names(&dir.join(included.trim()), visited)?);
143 continue;
144 }
145
146 if line.starts_with('-') {
147 if line.starts_with("-e") || line.starts_with("--editable") {
150 return None;
151 }
152 continue;
153 }
154
155 let spec = line.split(" @ ").next().unwrap_or(line).trim();
157 names.insert(requirement_name(spec)?);
158 }
159
160 Some(names)
161}
162
163pub(super) fn installed_distributions(venv: &Path) -> Option<HashMap<String, Vec<String>>> {
170 let mut site_packages: Vec<PathBuf> = Vec::new();
171 let windows_layout = venv.join("Lib").join("site-packages");
172 if windows_layout.is_dir() {
173 site_packages.push(windows_layout);
174 }
175 for lib in ["lib", "lib64"] {
178 let Ok(entries) = fs::read_dir(venv.join(lib)) else {
179 continue;
180 };
181 for entry in entries.flatten() {
182 let sp = entry.path().join("site-packages");
183 if sp.is_dir() {
184 site_packages.push(sp);
185 }
186 }
187 }
188 if site_packages.is_empty() {
189 return None;
190 }
191
192 let mut installed = HashMap::new();
193 for sp in site_packages {
194 let Ok(entries) = fs::read_dir(&sp) else {
195 continue;
196 };
197 for entry in entries.flatten() {
198 let file_name = entry.file_name().to_string_lossy().into_owned();
199 let Some(stem) = file_name
200 .strip_suffix(".dist-info")
201 .or_else(|| file_name.strip_suffix(".egg-info"))
202 else {
203 continue;
204 };
205 let name = stem.rsplit_once('-').map(|(n, _)| n).unwrap_or(stem);
207 installed.insert(
208 normalize_package_name(name),
209 declared_dependencies(&entry.path()),
210 );
211 }
212 }
213 Some(installed)
214}
215
216fn declared_dependencies(dist_info: &Path) -> Vec<String> {
221 let Ok(metadata) = fs::read_to_string(dist_info.join("METADATA")) else {
222 return Vec::new();
223 };
224 let mut deps = Vec::new();
225 for line in metadata.lines() {
226 if line.is_empty() {
229 break;
230 }
231 if let Some(spec) = line.strip_prefix("Requires-Dist:") {
232 if let Some(name) = requirement_name(spec.trim()) {
233 deps.push(name);
234 }
235 }
236 }
237 deps
238}
239
240fn venv_python_version(venv: &Path) -> Option<(u64, u64)> {
242 let cfg = fs::read_to_string(venv.join(PYVENV_CFG)).ok()?;
243 for line in cfg.lines() {
244 let Some((key, value)) = line.split_once('=') else {
245 continue;
246 };
247 if matches!(key.trim(), "version" | "version_info") {
248 let mut parts = value.trim().split('.');
249 return Some((parts.next()?.parse().ok()?, parts.next()?.parse().ok()?));
250 }
251 }
252 None
253}
254
255fn path_python_version() -> Option<(u64, u64)> {
258 let output = std::process::Command::new(super::resolve_program("python"))
259 .arg("--version")
260 .stdin(std::process::Stdio::null())
261 .output()
262 .ok()?;
263 if !output.status.success() {
264 return None;
265 }
266 let stdout = String::from_utf8_lossy(&output.stdout).trim().to_string();
268 let text = if stdout.is_empty() {
269 String::from_utf8_lossy(&output.stderr).trim().to_string()
270 } else {
271 stdout
272 };
273 let version = text.split_whitespace().nth(1)?;
274 let mut parts = version.split('.');
275 Some((parts.next()?.parse().ok()?, parts.next()?.parse().ok()?))
276}
277
278fn warn_about_restore_surprises(path: &Path, venvs: &[PathBuf]) {
282 if venvs.len() > 1 {
283 crate::output::print_warning(&format!(
284 "`{}` has {} virtual environments, all rebuilt from one requirements.txt. \
285 Each restores under its own recorded name; a plain `devp restore` with no \
286 record rebuilds only `.venv`.",
287 path.display(),
288 venvs.len()
289 ));
290 } else if let Some(venv) = venvs.first() {
291 let name = venv.file_name().map(|n| n.to_string_lossy().into_owned());
292 if let Some(name) = name {
293 if name != ".venv" {
294 crate::output::print_info(&format!(
295 "The environment at `{}` is named `{name}` — `devp restore --last-run` \
296 recreates that name, but a restore with no record creates `.venv`.",
297 venv.display()
298 ));
299 }
300 }
301 }
302
303 let on_path = path_python_version();
304 for venv in venvs {
305 if let (Some(built_with), Some(available)) = (venv_python_version(venv), on_path) {
306 if built_with != available {
307 crate::output::print_warning(&format!(
308 "`{}` was built with Python {}.{}, but `python` on PATH is {}.{} — a \
309 restore would rebuild it on that interpreter instead, and pinned \
310 wheels may not exist for it.",
311 venv.display(),
312 built_with.0,
313 built_with.1,
314 available.0,
315 available.1
316 ));
317 }
318 }
319 }
320}
321
322fn unrecorded_packages(
330 installed: &HashMap<String, Vec<String>>,
331 pinned: &HashSet<String>,
332) -> Vec<String> {
333 let mut reachable: HashSet<String> = HashSet::new();
334 let mut queue: Vec<String> = pinned.iter().cloned().collect();
335 queue.extend(BASELINE_DISTRIBUTIONS.iter().map(|s| (*s).to_string()));
336
337 while let Some(name) = queue.pop() {
338 if !reachable.insert(name.clone()) {
339 continue;
340 }
341 if let Some(deps) = installed.get(&name) {
342 queue.extend(deps.iter().cloned());
343 }
344 }
345
346 let mut extras: Vec<String> = installed
347 .keys()
348 .filter(|name| !reachable.contains(*name))
349 .cloned()
350 .collect();
351 extras.sort();
352 extras
353}
354
355impl PackageManager for Venv {
356 fn name(&self) -> &'static str {
357 "venv"
358 }
359
360 fn detect(&self, path: &Path) -> bool {
368 let req_txt = path.join("requirements.txt");
369 let uv_lock = path.join("uv.lock");
370
371 if !req_txt.exists() || uv_lock.exists() {
372 return false;
373 }
374
375 if FOREIGN_PYTHON_LOCKFILES
379 .iter()
380 .any(|f| path.join(f).exists())
381 || is_poetry_project(path)
382 {
383 return false;
384 }
385
386 !find_venv_dirs(path).is_empty()
387 }
388
389 fn bloat_dirs(&self, path: &Path) -> Vec<BloatDir> {
391 find_venv_dirs(path)
392 .into_iter()
393 .map(|venv_path| {
394 let name = venv_path
395 .file_name()
396 .map(|n| n.to_string_lossy().to_string())
397 .unwrap_or_else(|| venv_path.display().to_string());
398 let size = dir_size(&venv_path);
399 BloatDir {
400 name,
401 path: venv_path,
402 size_bytes: size,
403 shared_bytes: 0,
404 }
405 })
406 .collect()
407 }
408
409 fn enforce_lockfile(&self, path: &Path, _policy: EnforcePolicy) -> Result<()> {
412 let req_txt = path.join("requirements.txt");
413 if !req_txt.exists() {
414 return Err(anyhow!("requirements.txt missing"));
415 }
416 let has_requirements = fs::read_to_string(&req_txt)
419 .map(|c| {
420 c.lines()
421 .any(|l| !l.trim().is_empty() && !l.trim_start().starts_with('#'))
422 })
423 .unwrap_or(false);
424 if !has_requirements {
425 return Err(anyhow!(
426 "requirements.txt at `{}` lists no packages — the virtual environment \
427 could not be rebuilt after deletion. Populate it with `pip freeze > requirements.txt`.",
428 req_txt.display()
429 ));
430 }
431
432 let venvs = find_venv_dirs(path);
433 warn_about_restore_surprises(path, &venvs);
434
435 if let Some(pinned) = requirement_names(&req_txt, &mut Vec::new()) {
441 for venv in venvs {
442 let Some(installed) = installed_distributions(&venv) else {
443 continue;
444 };
445 let extras = unrecorded_packages(&installed, &pinned);
446 if extras.is_empty() {
447 continue;
448 }
449 let shown = extras
450 .iter()
451 .take(10)
452 .cloned()
453 .collect::<Vec<_>>()
454 .join(", ");
455 let suffix = if extras.len() > 10 {
456 format!(", … and {} more", extras.len() - 10)
457 } else {
458 String::new()
459 };
460 return Err(anyhow!(
461 "`{}` holds {} package(s) that requirements.txt does not account for \
462 ({shown}{suffix}). Deleting the environment would lose them with no \
463 way back. Record them first: `pip freeze > requirements.txt`.",
464 venv.display(),
465 extras.len()
466 ));
467 }
468 }
469 Ok(())
470 }
471
472 fn restore(&self, path: &Path, timeout: std::time::Duration) -> Result<()> {
476 self.restore_named(path, ".venv", timeout)
477 }
478
479 fn restore_named(
482 &self,
483 path: &Path,
484 dir_name: &str,
485 timeout: std::time::Duration,
486 ) -> Result<()> {
487 let dir_name = if dir_name.is_empty()
490 || dir_name == "."
491 || dir_name == ".."
492 || dir_name.contains(['/', '\\'])
493 {
494 ".venv"
495 } else {
496 dir_name
497 };
498 run_command_with_timeout("python", &["-m", "venv", dir_name], path, timeout)?;
499 #[cfg(windows)]
502 let python = path.join(dir_name).join("Scripts").join("python.exe");
503 #[cfg(not(windows))]
504 let python = path.join(dir_name).join("bin").join("python");
505 run_command_with_timeout(
506 &python.to_string_lossy(),
507 &["-m", "pip", "install", "-r", "requirements.txt"],
508 path,
509 timeout,
510 )
511 }
512
513 fn lockfiles(&self) -> &'static [&'static str] {
517 &["requirements.txt"]
518 }
519
520 fn drift(&self, path: &Path) -> Vec<super::DriftReport> {
523 let Some(pinned) = requirement_names(&path.join("requirements.txt"), &mut Vec::new())
524 else {
525 return Vec::new();
526 };
527 let mut reports = Vec::new();
528 for venv in find_venv_dirs(path) {
529 let Some(installed) = installed_distributions(&venv) else {
530 continue;
531 };
532 let extras = unrecorded_packages(&installed, &pinned);
533 if extras.is_empty() {
534 continue;
535 }
536 reports.push(super::DriftReport {
537 directory: venv
538 .file_name()
539 .map(|n| n.to_string_lossy().into_owned())
540 .unwrap_or_else(|| venv.display().to_string()),
541 unrecorded: extras,
542 record_command: "pip freeze > requirements.txt",
543 });
544 }
545 reports
546 }
547}
548
549#[cfg(test)]
550mod tests {
551 use super::*;
552 use std::fs::{self, File};
553 use tempfile::tempdir;
554
555 fn make_venv(dir: &Path, name: &str) {
556 let venv = dir.join(name);
557 fs::create_dir(&venv).unwrap();
558 File::create(venv.join(PYVENV_CFG)).unwrap();
559 }
560
561 #[test]
562 fn test_name() {
563 assert_eq!(Venv.name(), "venv");
564 }
565
566 #[test]
567 fn test_detect_positive_dot_venv() {
568 let dir = tempdir().unwrap();
569 File::create(dir.path().join("requirements.txt")).unwrap();
570 make_venv(dir.path(), ".venv");
571 assert!(Venv.detect(dir.path()));
572 }
573
574 #[test]
575 fn test_detect_positive_venv() {
576 let dir = tempdir().unwrap();
577 File::create(dir.path().join("requirements.txt")).unwrap();
578 make_venv(dir.path(), "venv");
579 assert!(Venv.detect(dir.path()));
580 }
581
582 #[test]
583 fn test_detect_positive_custom_name() {
584 let dir = tempdir().unwrap();
585 File::create(dir.path().join("requirements.txt")).unwrap();
586 make_venv(dir.path(), "my_env");
587 assert!(Venv.detect(dir.path()));
588 }
589
590 #[test]
591 fn test_detect_positive_env() {
592 let dir = tempdir().unwrap();
593 File::create(dir.path().join("requirements.txt")).unwrap();
594 make_venv(dir.path(), "env");
595 assert!(Venv.detect(dir.path()));
596 }
597
598 #[test]
599 fn test_detect_negative_no_req() {
600 let dir = tempdir().unwrap();
601 make_venv(dir.path(), ".venv");
602 assert!(!Venv.detect(dir.path()));
603 }
604
605 #[test]
606 fn test_detect_negative_no_env() {
607 let dir = tempdir().unwrap();
608 File::create(dir.path().join("requirements.txt")).unwrap();
609 fs::create_dir(dir.path().join("not_a_venv")).unwrap();
611 assert!(!Venv.detect(dir.path()));
612 }
613
614 #[test]
615 fn test_detect_negative_uv_lock() {
616 let dir = tempdir().unwrap();
617 File::create(dir.path().join("requirements.txt")).unwrap();
618 File::create(dir.path().join("uv.lock")).unwrap();
619 make_venv(dir.path(), ".venv");
620 assert!(!Venv.detect(dir.path()));
621 }
622
623 #[test]
624 fn test_bloat_dirs_present() {
625 let dir = tempdir().unwrap();
626 make_venv(dir.path(), ".venv");
627 make_venv(dir.path(), "my_env");
628 let dirs = Venv.bloat_dirs(dir.path());
629 assert_eq!(dirs.len(), 2);
630 let names: Vec<&str> = dirs.iter().map(|d| d.name.as_str()).collect();
631 assert!(names.contains(&".venv"));
632 assert!(names.contains(&"my_env"));
633 }
634
635 #[test]
636 fn test_bloat_dirs_absent() {
637 let dir = tempdir().unwrap();
638 let dirs = Venv.bloat_dirs(dir.path());
639 assert!(dirs.is_empty());
640 }
641
642 #[test]
643 fn test_bloat_dirs_ignores_non_venv_dirs() {
644 let dir = tempdir().unwrap();
645 fs::create_dir(dir.path().join("src")).unwrap();
647 make_venv(dir.path(), ".venv");
648 let dirs = Venv.bloat_dirs(dir.path());
649 assert_eq!(dirs.len(), 1);
650 assert_eq!(dirs[0].name, ".venv");
651 }
652
653 fn install_package(root: &Path, venv: &str, name: &str, requires: &[&str]) {
657 let dist_info = root
658 .join(venv)
659 .join("Lib")
660 .join("site-packages")
661 .join(format!("{name}-1.0.0.dist-info"));
662 fs::create_dir_all(&dist_info).unwrap();
663 let mut metadata = format!("Metadata-Version: 2.1\nName: {name}\nVersion: 1.0.0\n");
664 for dep in requires {
665 metadata.push_str(&format!("Requires-Dist: {dep}\n"));
666 }
667 fs::write(dist_info.join("METADATA"), metadata).unwrap();
668 }
669
670 #[test]
671 fn enforce_refuses_when_requirements_lists_nothing() {
672 let dir = tempdir().unwrap();
673 fs::write(dir.path().join("requirements.txt"), "# nothing pinned\n\n").unwrap();
674 make_venv(dir.path(), ".venv");
675
676 let err = Venv
677 .enforce_lockfile(dir.path(), EnforcePolicy::default())
678 .unwrap_err();
679 assert!(err.to_string().contains("lists no packages"));
680 }
681
682 #[test]
683 fn enforce_refuses_a_package_the_requirements_never_recorded() {
684 let dir = tempdir().unwrap();
687 fs::write(dir.path().join("requirements.txt"), "flask==3.0.0\n").unwrap();
688 make_venv(dir.path(), ".venv");
689 install_package(dir.path(), ".venv", "flask", &[]);
690 install_package(dir.path(), ".venv", "requests", &[]);
691
692 let err = Venv
693 .enforce_lockfile(dir.path(), EnforcePolicy::default())
694 .unwrap_err()
695 .to_string();
696 assert!(
697 err.contains("requests"),
698 "names the unrecorded package: {err}"
699 );
700 assert!(
701 !err.contains("flask"),
702 "must not blame the pinned one: {err}"
703 );
704 assert!(err.contains("pip freeze"), "says how to record it: {err}");
705 }
706
707 #[test]
708 fn enforce_accepts_transitive_dependencies_of_pinned_packages() {
709 let dir = tempdir().unwrap();
712 fs::write(dir.path().join("requirements.txt"), "requests==2.32.3\n").unwrap();
713 make_venv(dir.path(), ".venv");
714 install_package(
715 dir.path(),
716 ".venv",
717 "requests",
718 &["urllib3 (>=1.21.1)", "charset-normalizer"],
719 );
720 install_package(dir.path(), ".venv", "urllib3", &[]);
721 install_package(dir.path(), ".venv", "charset_normalizer", &[]);
723
724 assert!(
725 Venv.enforce_lockfile(dir.path(), EnforcePolicy::default())
726 .is_ok()
727 );
728 }
729
730 #[test]
731 fn enforce_skips_the_comparison_when_requirements_cannot_be_parsed() {
732 let dir = tempdir().unwrap();
735 fs::write(
736 dir.path().join("requirements.txt"),
737 "-e ./local-package\nflask==3.0.0\n",
738 )
739 .unwrap();
740 make_venv(dir.path(), ".venv");
741 install_package(dir.path(), ".venv", "left-behind", &[]);
742
743 assert!(
744 Venv.enforce_lockfile(dir.path(), EnforcePolicy::default())
745 .is_ok()
746 );
747 }
748
749 #[test]
750 fn drift_names_the_venv_and_the_unrecorded_packages() {
751 let dir = tempdir().unwrap();
752 fs::write(dir.path().join("requirements.txt"), "requests==2.32.3\n").unwrap();
753 make_venv(dir.path(), ".venv");
754 install_package(dir.path(), ".venv", "requests", &[]);
755 install_package(dir.path(), ".venv", "sneaky-pkg", &[]);
756
757 let reports = Venv.drift(dir.path());
758 assert_eq!(reports.len(), 1);
759 assert_eq!(reports[0].directory, ".venv");
760 assert_eq!(reports[0].unrecorded, vec!["sneaky-pkg"]);
761 assert_eq!(reports[0].record_command, "pip freeze > requirements.txt");
762 }
763
764 #[test]
767 fn drift_does_not_flag_transitive_dependencies_of_pinned_packages() {
768 let dir = tempdir().unwrap();
769 fs::write(dir.path().join("requirements.txt"), "requests==2.32.3\n").unwrap();
770 make_venv(dir.path(), ".venv");
771 install_package(dir.path(), ".venv", "requests", &["urllib3"]);
772 install_package(dir.path(), ".venv", "urllib3", &[]);
773
774 assert!(Venv.drift(dir.path()).is_empty());
775 }
776}