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
212 .match_indices('-')
213 .find(|(i, _)| {
214 stem[i + 1..]
215 .chars()
216 .next()
217 .is_some_and(|c| c.is_ascii_digit())
218 })
219 .map(|(i, _)| &stem[..i])
220 .unwrap_or(stem);
221 installed.insert(
222 normalize_package_name(name),
223 declared_dependencies(&entry.path()),
224 );
225 }
226 }
227 Some(installed)
228}
229
230fn declared_dependencies(dist_info: &Path) -> Vec<String> {
235 let Ok(metadata) = fs::read_to_string(dist_info.join("METADATA")) else {
236 return Vec::new();
237 };
238 let mut deps = Vec::new();
239 for line in metadata.lines() {
240 if line.is_empty() {
243 break;
244 }
245 if let Some(spec) = line.strip_prefix("Requires-Dist:")
246 && let Some(name) = requirement_name(spec.trim())
247 {
248 deps.push(name);
249 }
250 }
251 deps
252}
253
254fn venv_python_version(venv: &Path) -> Option<(u64, u64)> {
256 let cfg = fs::read_to_string(venv.join(PYVENV_CFG)).ok()?;
257 for line in cfg.lines() {
258 let Some((key, value)) = line.split_once('=') else {
259 continue;
260 };
261 if matches!(key.trim(), "version" | "version_info") {
262 let mut parts = value.trim().split('.');
263 return Some((parts.next()?.parse().ok()?, parts.next()?.parse().ok()?));
264 }
265 }
266 None
267}
268
269fn path_python_version() -> Option<(u64, u64)> {
272 let output = crate::spawn::command(super::resolve_program("python"))
273 .arg("--version")
274 .stdin(std::process::Stdio::null())
275 .output()
276 .ok()?;
277 if !output.status.success() {
278 return None;
279 }
280 let stdout = String::from_utf8_lossy(&output.stdout).trim().to_string();
282 let text = if stdout.is_empty() {
283 String::from_utf8_lossy(&output.stderr).trim().to_string()
284 } else {
285 stdout
286 };
287 let version = text.split_whitespace().nth(1)?;
288 let mut parts = version.split('.');
289 Some((parts.next()?.parse().ok()?, parts.next()?.parse().ok()?))
290}
291
292fn warn_about_restore_surprises(path: &Path, venvs: &[PathBuf]) {
296 if venvs.len() > 1 {
297 crate::output::print_warning(&format!(
298 "`{}` has {} virtual environments, all rebuilt from one requirements.txt. \
299 Each restores under its own recorded name; a plain `devp restore` with no \
300 record rebuilds only `.venv`.",
301 path.display(),
302 venvs.len()
303 ));
304 } else if let Some(venv) = venvs.first() {
305 let name = venv.file_name().map(|n| n.to_string_lossy().into_owned());
306 if let Some(name) = name
307 && name != ".venv"
308 {
309 crate::output::print_info(&format!(
310 "The environment at `{}` is named `{name}` — `devp restore --last-run` \
311 recreates that name, but a restore with no record creates `.venv`.",
312 venv.display()
313 ));
314 }
315 }
316
317 let on_path = path_python_version();
318 for venv in venvs {
319 if let (Some(built_with), Some(available)) = (venv_python_version(venv), on_path)
320 && built_with != available
321 {
322 crate::output::print_warning(&format!(
323 "`{}` was built with Python {}.{}, but `python` on PATH is {}.{} — a \
324 restore would rebuild it on that interpreter instead, and pinned \
325 wheels may not exist for it.",
326 venv.display(),
327 built_with.0,
328 built_with.1,
329 available.0,
330 available.1
331 ));
332 }
333 }
334}
335
336fn unrecorded_packages(
344 installed: &HashMap<String, Vec<String>>,
345 pinned: &HashSet<String>,
346) -> Vec<String> {
347 let mut reachable: HashSet<String> = HashSet::new();
348 let mut queue: Vec<String> = pinned.iter().cloned().collect();
349 queue.extend(BASELINE_DISTRIBUTIONS.iter().map(|s| (*s).to_string()));
350
351 while let Some(name) = queue.pop() {
352 if !reachable.insert(name.clone()) {
353 continue;
354 }
355 if let Some(deps) = installed.get(&name) {
356 queue.extend(deps.iter().cloned());
357 }
358 }
359
360 let mut extras: Vec<String> = installed
361 .keys()
362 .filter(|name| !reachable.contains(*name))
363 .cloned()
364 .collect();
365 extras.sort();
366 extras
367}
368
369impl PackageManager for Venv {
370 fn name(&self) -> &'static str {
371 "venv"
372 }
373
374 fn detect(&self, path: &Path) -> bool {
382 let req_txt = path.join("requirements.txt");
383 let uv_lock = path.join("uv.lock");
384
385 if !req_txt.exists() || uv_lock.exists() {
386 return false;
387 }
388
389 if FOREIGN_PYTHON_LOCKFILES
393 .iter()
394 .any(|f| path.join(f).exists())
395 || is_poetry_project(path)
396 {
397 return false;
398 }
399
400 !find_venv_dirs(path).is_empty()
401 }
402
403 fn bloat_dirs(&self, path: &Path) -> Vec<BloatDir> {
405 find_venv_dirs(path)
406 .into_iter()
407 .map(|venv_path| {
408 let name = venv_path
409 .file_name()
410 .map(|n| n.to_string_lossy().to_string())
411 .unwrap_or_else(|| venv_path.display().to_string());
412 let size = dir_size(&venv_path);
413 BloatDir {
414 name,
415 path: venv_path,
416 size_bytes: size,
417 shared_bytes: 0,
418 }
419 })
420 .collect()
421 }
422
423 fn enforce_lockfile(&self, path: &Path, _policy: EnforcePolicy) -> Result<()> {
426 let req_txt = path.join("requirements.txt");
427 if !req_txt.exists() {
428 return Err(anyhow!("requirements.txt missing"));
429 }
430 let has_requirements = fs::read_to_string(&req_txt)
433 .map(|c| {
434 c.lines()
435 .any(|l| !l.trim().is_empty() && !l.trim_start().starts_with('#'))
436 })
437 .unwrap_or(false);
438 if !has_requirements {
439 return Err(anyhow!(
440 "requirements.txt at `{}` lists no packages — the virtual environment \
441 could not be rebuilt after deletion. Populate it with `pip freeze > requirements.txt`.",
442 req_txt.display()
443 ));
444 }
445
446 let venvs = find_venv_dirs(path);
447 warn_about_restore_surprises(path, &venvs);
448
449 if let Some(pinned) = requirement_names(&req_txt, &mut Vec::new()) {
455 for venv in venvs {
456 let Some(installed) = installed_distributions(&venv) else {
457 continue;
458 };
459 let extras = unrecorded_packages(&installed, &pinned);
460 if extras.is_empty() {
461 continue;
462 }
463 let shown = extras
464 .iter()
465 .take(10)
466 .cloned()
467 .collect::<Vec<_>>()
468 .join(", ");
469 let suffix = if extras.len() > 10 {
470 format!(", … and {} more", extras.len() - 10)
471 } else {
472 String::new()
473 };
474 return Err(anyhow!(
475 "`{}` holds {} package(s) that requirements.txt does not account for \
476 ({shown}{suffix}). Deleting the environment would lose them with no \
477 way back. Record them first: `pip freeze > requirements.txt`.",
478 venv.display(),
479 extras.len()
480 ));
481 }
482 }
483 Ok(())
484 }
485
486 fn restore(&self, path: &Path, timeout: std::time::Duration) -> Result<()> {
490 self.restore_named(path, ".venv", timeout)
491 }
492
493 fn restore_named(
496 &self,
497 path: &Path,
498 dir_name: &str,
499 timeout: std::time::Duration,
500 ) -> Result<()> {
501 let dir_name = if dir_name.is_empty()
504 || dir_name == "."
505 || dir_name == ".."
506 || dir_name.contains(['/', '\\'])
507 {
508 ".venv"
509 } else {
510 dir_name
511 };
512 run_command_with_timeout("python", &["-m", "venv", dir_name], path, timeout)?;
513 #[cfg(windows)]
516 let python = path.join(dir_name).join("Scripts").join("python.exe");
517 #[cfg(not(windows))]
518 let python = path.join(dir_name).join("bin").join("python");
519 run_command_with_timeout(
520 &python.to_string_lossy(),
521 &["-m", "pip", "install", "-r", "requirements.txt"],
522 path,
523 timeout,
524 )
525 }
526
527 fn lockfiles(&self) -> &'static [&'static str] {
531 &["requirements.txt"]
532 }
533
534 fn drift(&self, path: &Path) -> Vec<super::DriftReport> {
537 let Some(pinned) = requirement_names(&path.join("requirements.txt"), &mut Vec::new())
538 else {
539 return Vec::new();
540 };
541 let mut reports = Vec::new();
542 for venv in find_venv_dirs(path) {
543 let Some(installed) = installed_distributions(&venv) else {
544 continue;
545 };
546 let extras = unrecorded_packages(&installed, &pinned);
547 if extras.is_empty() {
548 continue;
549 }
550 reports.push(super::DriftReport {
551 directory: venv
552 .file_name()
553 .map(|n| n.to_string_lossy().into_owned())
554 .unwrap_or_else(|| venv.display().to_string()),
555 unrecorded: extras,
556 record_command: "pip freeze > requirements.txt",
557 });
558 }
559 reports
560 }
561}
562
563#[cfg(test)]
564mod tests {
565 use super::*;
566 use std::fs::{self, File};
567 use tempfile::tempdir;
568
569 fn make_venv(dir: &Path, name: &str) {
570 let venv = dir.join(name);
571 fs::create_dir(&venv).unwrap();
572 File::create(venv.join(PYVENV_CFG)).unwrap();
573 }
574
575 #[test]
576 fn test_name() {
577 assert_eq!(Venv.name(), "venv");
578 }
579
580 #[test]
581 fn test_detect_positive_dot_venv() {
582 let dir = tempdir().unwrap();
583 File::create(dir.path().join("requirements.txt")).unwrap();
584 make_venv(dir.path(), ".venv");
585 assert!(Venv.detect(dir.path()));
586 }
587
588 #[test]
589 fn test_detect_positive_venv() {
590 let dir = tempdir().unwrap();
591 File::create(dir.path().join("requirements.txt")).unwrap();
592 make_venv(dir.path(), "venv");
593 assert!(Venv.detect(dir.path()));
594 }
595
596 #[test]
597 fn test_detect_positive_custom_name() {
598 let dir = tempdir().unwrap();
599 File::create(dir.path().join("requirements.txt")).unwrap();
600 make_venv(dir.path(), "my_env");
601 assert!(Venv.detect(dir.path()));
602 }
603
604 #[test]
605 fn test_detect_positive_env() {
606 let dir = tempdir().unwrap();
607 File::create(dir.path().join("requirements.txt")).unwrap();
608 make_venv(dir.path(), "env");
609 assert!(Venv.detect(dir.path()));
610 }
611
612 #[test]
613 fn test_detect_negative_no_req() {
614 let dir = tempdir().unwrap();
615 make_venv(dir.path(), ".venv");
616 assert!(!Venv.detect(dir.path()));
617 }
618
619 #[test]
620 fn test_detect_negative_no_env() {
621 let dir = tempdir().unwrap();
622 File::create(dir.path().join("requirements.txt")).unwrap();
623 fs::create_dir(dir.path().join("not_a_venv")).unwrap();
625 assert!(!Venv.detect(dir.path()));
626 }
627
628 #[test]
629 fn test_detect_negative_uv_lock() {
630 let dir = tempdir().unwrap();
631 File::create(dir.path().join("requirements.txt")).unwrap();
632 File::create(dir.path().join("uv.lock")).unwrap();
633 make_venv(dir.path(), ".venv");
634 assert!(!Venv.detect(dir.path()));
635 }
636
637 #[test]
638 fn test_bloat_dirs_present() {
639 let dir = tempdir().unwrap();
640 make_venv(dir.path(), ".venv");
641 make_venv(dir.path(), "my_env");
642 let dirs = Venv.bloat_dirs(dir.path());
643 assert_eq!(dirs.len(), 2);
644 let names: Vec<&str> = dirs.iter().map(|d| d.name.as_str()).collect();
645 assert!(names.contains(&".venv"));
646 assert!(names.contains(&"my_env"));
647 }
648
649 #[test]
650 fn test_bloat_dirs_absent() {
651 let dir = tempdir().unwrap();
652 let dirs = Venv.bloat_dirs(dir.path());
653 assert!(dirs.is_empty());
654 }
655
656 #[test]
657 fn test_bloat_dirs_ignores_non_venv_dirs() {
658 let dir = tempdir().unwrap();
659 fs::create_dir(dir.path().join("src")).unwrap();
661 make_venv(dir.path(), ".venv");
662 let dirs = Venv.bloat_dirs(dir.path());
663 assert_eq!(dirs.len(), 1);
664 assert_eq!(dirs[0].name, ".venv");
665 }
666
667 fn install_package(root: &Path, venv: &str, name: &str, requires: &[&str]) {
671 let dist_info = root
672 .join(venv)
673 .join("Lib")
674 .join("site-packages")
675 .join(format!("{name}-1.0.0.dist-info"));
676 fs::create_dir_all(&dist_info).unwrap();
677 let mut metadata = format!("Metadata-Version: 2.1\nName: {name}\nVersion: 1.0.0\n");
678 for dep in requires {
679 metadata.push_str(&format!("Requires-Dist: {dep}\n"));
680 }
681 fs::write(dist_info.join("METADATA"), metadata).unwrap();
682 }
683
684 #[test]
685 fn enforce_refuses_when_requirements_lists_nothing() {
686 let dir = tempdir().unwrap();
687 fs::write(dir.path().join("requirements.txt"), "# nothing pinned\n\n").unwrap();
688 make_venv(dir.path(), ".venv");
689
690 let err = Venv
691 .enforce_lockfile(dir.path(), EnforcePolicy::default())
692 .unwrap_err();
693 assert!(err.to_string().contains("lists no packages"));
694 }
695
696 #[test]
697 fn enforce_refuses_a_package_the_requirements_never_recorded() {
698 let dir = tempdir().unwrap();
701 fs::write(dir.path().join("requirements.txt"), "flask==3.0.0\n").unwrap();
702 make_venv(dir.path(), ".venv");
703 install_package(dir.path(), ".venv", "flask", &[]);
704 install_package(dir.path(), ".venv", "requests", &[]);
705
706 let err = Venv
707 .enforce_lockfile(dir.path(), EnforcePolicy::default())
708 .unwrap_err()
709 .to_string();
710 assert!(
711 err.contains("requests"),
712 "names the unrecorded package: {err}"
713 );
714 assert!(
715 !err.contains("flask"),
716 "must not blame the pinned one: {err}"
717 );
718 assert!(err.contains("pip freeze"), "says how to record it: {err}");
719 }
720
721 #[test]
722 fn enforce_accepts_transitive_dependencies_of_pinned_packages() {
723 let dir = tempdir().unwrap();
726 fs::write(dir.path().join("requirements.txt"), "requests==2.32.3\n").unwrap();
727 make_venv(dir.path(), ".venv");
728 install_package(
729 dir.path(),
730 ".venv",
731 "requests",
732 &["urllib3 (>=1.21.1)", "charset-normalizer"],
733 );
734 install_package(dir.path(), ".venv", "urllib3", &[]);
735 install_package(dir.path(), ".venv", "charset_normalizer", &[]);
737
738 assert!(
739 Venv.enforce_lockfile(dir.path(), EnforcePolicy::default())
740 .is_ok()
741 );
742 }
743
744 #[test]
745 fn enforce_skips_the_comparison_when_requirements_cannot_be_parsed() {
746 let dir = tempdir().unwrap();
749 fs::write(
750 dir.path().join("requirements.txt"),
751 "-e ./local-package\nflask==3.0.0\n",
752 )
753 .unwrap();
754 make_venv(dir.path(), ".venv");
755 install_package(dir.path(), ".venv", "left-behind", &[]);
756
757 assert!(
758 Venv.enforce_lockfile(dir.path(), EnforcePolicy::default())
759 .is_ok()
760 );
761 }
762
763 #[test]
764 fn drift_names_the_venv_and_the_unrecorded_packages() {
765 let dir = tempdir().unwrap();
766 fs::write(dir.path().join("requirements.txt"), "requests==2.32.3\n").unwrap();
767 make_venv(dir.path(), ".venv");
768 install_package(dir.path(), ".venv", "requests", &[]);
769 install_package(dir.path(), ".venv", "sneaky-pkg", &[]);
770
771 let reports = Venv.drift(dir.path());
772 assert_eq!(reports.len(), 1);
773 assert_eq!(reports[0].directory, ".venv");
774 assert_eq!(reports[0].unrecorded, vec!["sneaky-pkg"]);
775 assert_eq!(reports[0].record_command, "pip freeze > requirements.txt");
776 }
777
778 #[test]
781 fn drift_does_not_flag_transitive_dependencies_of_pinned_packages() {
782 let dir = tempdir().unwrap();
783 fs::write(dir.path().join("requirements.txt"), "requests==2.32.3\n").unwrap();
784 make_venv(dir.path(), ".venv");
785 install_package(dir.path(), ".venv", "requests", &["urllib3"]);
786 install_package(dir.path(), ".venv", "urllib3", &[]);
787
788 assert!(Venv.drift(dir.path()).is_empty());
789 }
790}