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 tag = super::venv_runtime_tag(venv)?;
257 let (major, minor) = tag.split_once('.')?;
258 Some((major.parse().ok()?, minor.parse().ok()?))
259}
260
261fn path_python_version() -> Option<(u64, u64)> {
264 let output = crate::spawn::command(super::resolve_program("python"))
265 .arg("--version")
266 .stdin(std::process::Stdio::null())
267 .output()
268 .ok()?;
269 if !output.status.success() {
270 return None;
271 }
272 let stdout = String::from_utf8_lossy(&output.stdout).trim().to_string();
274 let text = if stdout.is_empty() {
275 String::from_utf8_lossy(&output.stderr).trim().to_string()
276 } else {
277 stdout
278 };
279 let version = text.split_whitespace().nth(1)?;
280 let mut parts = version.split('.');
281 Some((parts.next()?.parse().ok()?, parts.next()?.parse().ok()?))
282}
283
284fn warn_about_restore_surprises(path: &Path, venvs: &[PathBuf]) {
288 if venvs.len() > 1 {
289 crate::output::print_warning(&format!(
290 "`{}` has {} virtual environments, all rebuilt from one requirements.txt. \
291 Each restores under its own recorded name; a plain `devp restore` with no \
292 record rebuilds only `.venv`.",
293 crate::output::clean_path(path),
294 venvs.len()
295 ));
296 } else if let Some(venv) = venvs.first() {
297 let name = venv.file_name().map(|n| n.to_string_lossy().into_owned());
298 if let Some(name) = name
299 && name != ".venv"
300 {
301 crate::output::print_info(&format!(
302 "The environment at `{}` is named `{name}` — `devp restore --last-run` \
303 recreates that name, but a restore with no record creates `.venv`.",
304 crate::output::clean_path(venv)
305 ));
306 }
307 }
308
309 let on_path = path_python_version();
310 for venv in venvs {
311 if let (Some(built_with), Some(available)) = (venv_python_version(venv), on_path)
312 && built_with != available
313 {
314 crate::output::print_warning(&format!(
315 "`{}` was built with Python {}.{}, but `python` on PATH is {}.{} — a \
316 restore would rebuild it on that interpreter instead, and pinned \
317 wheels may not exist for it.",
318 crate::output::clean_path(venv),
319 built_with.0,
320 built_with.1,
321 available.0,
322 available.1
323 ));
324 let dir = crate::output::clean_path(venv);
329 let (major, minor) = built_with;
330 #[cfg(windows)]
331 let native = format!("py -{major}.{minor} -m venv \"{dir}\"");
332 #[cfg(not(windows))]
333 let native = format!("python{major}.{minor} -m venv \"{dir}\"");
334 crate::output::print_info(&format!(
335 " Rebuild on {major}.{minor}: uv venv --python {major}.{minor} \"{dir}\" (or `{native}`)"
336 ));
337 }
338 }
339}
340
341fn unrecorded_packages(
349 installed: &HashMap<String, Vec<String>>,
350 pinned: &HashSet<String>,
351) -> Vec<String> {
352 let mut reachable: HashSet<String> = HashSet::new();
353 let mut queue: Vec<String> = pinned.iter().cloned().collect();
354 queue.extend(BASELINE_DISTRIBUTIONS.iter().map(|s| (*s).to_string()));
355
356 while let Some(name) = queue.pop() {
357 if !reachable.insert(name.clone()) {
358 continue;
359 }
360 if let Some(deps) = installed.get(&name) {
361 queue.extend(deps.iter().cloned());
362 }
363 }
364
365 let mut extras: Vec<String> = installed
366 .keys()
367 .filter(|name| !reachable.contains(*name))
368 .cloned()
369 .collect();
370 extras.sort();
371 extras
372}
373
374impl PackageManager for Venv {
375 fn name(&self) -> &'static str {
376 "venv"
377 }
378
379 fn detect(&self, path: &Path) -> bool {
387 let req_txt = path.join("requirements.txt");
388 let uv_lock = path.join("uv.lock");
389
390 if !req_txt.exists() || uv_lock.exists() {
391 return false;
392 }
393
394 if FOREIGN_PYTHON_LOCKFILES
398 .iter()
399 .any(|f| path.join(f).exists())
400 || is_poetry_project(path)
401 {
402 return false;
403 }
404
405 !find_venv_dirs(path).is_empty()
406 }
407
408 fn bloat_dirs(&self, path: &Path) -> Vec<BloatDir> {
410 find_venv_dirs(path)
411 .into_iter()
412 .map(|venv_path| {
413 let name = venv_path
414 .file_name()
415 .map(|n| n.to_string_lossy().to_string())
416 .unwrap_or_else(|| venv_path.display().to_string());
417 let size = dir_size(&venv_path);
418 BloatDir {
419 name,
420 path: venv_path,
421 size_bytes: size,
422 shared_bytes: 0,
423 }
424 })
425 .collect()
426 }
427
428 fn enforce_lockfile(&self, path: &Path, _policy: EnforcePolicy) -> Result<()> {
431 let req_txt = path.join("requirements.txt");
432 if !req_txt.exists() {
433 return Err(anyhow!("requirements.txt missing"));
434 }
435 let has_requirements = fs::read_to_string(&req_txt)
438 .map(|c| {
439 c.lines()
440 .any(|l| !l.trim().is_empty() && !l.trim_start().starts_with('#'))
441 })
442 .unwrap_or(false);
443 if !has_requirements {
444 return Err(anyhow!(
445 "requirements.txt at `{}` lists no packages — the virtual environment \
446 could not be rebuilt after deletion. Populate it with `pip freeze > requirements.txt`.",
447 req_txt.display()
448 ));
449 }
450
451 let venvs = find_venv_dirs(path);
452 warn_about_restore_surprises(path, &venvs);
453
454 if let Some(pinned) = requirement_names(&req_txt, &mut Vec::new()) {
460 for venv in venvs {
461 let Some(installed) = installed_distributions(&venv) else {
462 continue;
463 };
464 let extras = unrecorded_packages(&installed, &pinned);
465 if extras.is_empty() {
466 continue;
467 }
468 let shown = extras
469 .iter()
470 .take(10)
471 .cloned()
472 .collect::<Vec<_>>()
473 .join(", ");
474 let suffix = if extras.len() > 10 {
475 format!(", … and {} more", extras.len() - 10)
476 } else {
477 String::new()
478 };
479 return Err(anyhow!(
480 "`{}` holds {} package(s) that requirements.txt does not account for \
481 ({shown}{suffix}). Deleting the environment would lose them with no \
482 way back. Record them first: `pip freeze > requirements.txt`.",
483 venv.display(),
484 extras.len()
485 ));
486 }
487 }
488 Ok(())
489 }
490
491 fn restore(&self, path: &Path, timeout: std::time::Duration) -> Result<()> {
495 self.restore_named(path, ".venv", None, timeout)
496 }
497
498 fn runtime_tag(&self, path: &Path, dir_name: &str) -> Option<String> {
500 super::venv_runtime_tag(&path.join(dir_name))
501 }
502
503 fn restore_named(
506 &self,
507 path: &Path,
508 dir_name: &str,
509 runtime: Option<&str>,
510 timeout: std::time::Duration,
511 ) -> Result<()> {
512 let dir_name = if dir_name.is_empty()
515 || dir_name == "."
516 || dir_name == ".."
517 || dir_name.contains(['/', '\\'])
518 {
519 ".venv"
520 } else {
521 dir_name
522 };
523 let launcher = runtime
530 .filter(|tag| super::python_runtime_available(tag))
531 .and_then(super::python_launcher);
532 match launcher {
533 Some((program, prefix)) => {
534 let mut args: Vec<&str> = prefix.iter().map(String::as_str).collect();
535 args.extend_from_slice(&["-m", "venv", dir_name]);
536 run_command_with_timeout(&program, &args, path, timeout)?;
537 }
538 None => {
539 run_command_with_timeout("python", &["-m", "venv", dir_name], path, timeout)?;
540 }
541 }
542 #[cfg(windows)]
545 let python = path.join(dir_name).join("Scripts").join("python.exe");
546 #[cfg(not(windows))]
547 let python = path.join(dir_name).join("bin").join("python");
548 run_command_with_timeout(
549 &python.to_string_lossy(),
550 &["-m", "pip", "install", "-r", "requirements.txt"],
551 path,
552 timeout,
553 )
554 }
555
556 fn lockfiles(&self) -> &'static [&'static str] {
560 &["requirements.txt"]
561 }
562
563 fn drift(&self, path: &Path) -> Vec<super::DriftReport> {
566 let Some(pinned) = requirement_names(&path.join("requirements.txt"), &mut Vec::new())
567 else {
568 return Vec::new();
569 };
570 let mut reports = Vec::new();
571 for venv in find_venv_dirs(path) {
572 let Some(installed) = installed_distributions(&venv) else {
573 continue;
574 };
575 let extras = unrecorded_packages(&installed, &pinned);
576 if extras.is_empty() {
577 continue;
578 }
579 reports.push(super::DriftReport {
580 directory: venv
581 .file_name()
582 .map(|n| n.to_string_lossy().into_owned())
583 .unwrap_or_else(|| venv.display().to_string()),
584 unrecorded: extras,
585 record_command: "pip freeze > requirements.txt",
586 });
587 }
588 reports
589 }
590}
591
592#[cfg(test)]
593mod tests {
594 use super::*;
595 use std::fs::{self, File};
596 use tempfile::tempdir;
597
598 fn make_venv(dir: &Path, name: &str) {
599 let venv = dir.join(name);
600 fs::create_dir(&venv).unwrap();
601 File::create(venv.join(PYVENV_CFG)).unwrap();
602 }
603
604 #[test]
605 fn test_name() {
606 assert_eq!(Venv.name(), "venv");
607 }
608
609 #[test]
610 fn test_detect_positive_dot_venv() {
611 let dir = tempdir().unwrap();
612 File::create(dir.path().join("requirements.txt")).unwrap();
613 make_venv(dir.path(), ".venv");
614 assert!(Venv.detect(dir.path()));
615 }
616
617 #[test]
618 fn test_detect_positive_venv() {
619 let dir = tempdir().unwrap();
620 File::create(dir.path().join("requirements.txt")).unwrap();
621 make_venv(dir.path(), "venv");
622 assert!(Venv.detect(dir.path()));
623 }
624
625 #[test]
626 fn test_detect_positive_custom_name() {
627 let dir = tempdir().unwrap();
628 File::create(dir.path().join("requirements.txt")).unwrap();
629 make_venv(dir.path(), "my_env");
630 assert!(Venv.detect(dir.path()));
631 }
632
633 #[test]
634 fn test_detect_positive_env() {
635 let dir = tempdir().unwrap();
636 File::create(dir.path().join("requirements.txt")).unwrap();
637 make_venv(dir.path(), "env");
638 assert!(Venv.detect(dir.path()));
639 }
640
641 #[test]
642 fn test_detect_negative_no_req() {
643 let dir = tempdir().unwrap();
644 make_venv(dir.path(), ".venv");
645 assert!(!Venv.detect(dir.path()));
646 }
647
648 #[test]
649 fn test_detect_negative_no_env() {
650 let dir = tempdir().unwrap();
651 File::create(dir.path().join("requirements.txt")).unwrap();
652 fs::create_dir(dir.path().join("not_a_venv")).unwrap();
654 assert!(!Venv.detect(dir.path()));
655 }
656
657 #[test]
658 fn test_detect_negative_uv_lock() {
659 let dir = tempdir().unwrap();
660 File::create(dir.path().join("requirements.txt")).unwrap();
661 File::create(dir.path().join("uv.lock")).unwrap();
662 make_venv(dir.path(), ".venv");
663 assert!(!Venv.detect(dir.path()));
664 }
665
666 #[test]
667 fn test_bloat_dirs_present() {
668 let dir = tempdir().unwrap();
669 make_venv(dir.path(), ".venv");
670 make_venv(dir.path(), "my_env");
671 let dirs = Venv.bloat_dirs(dir.path());
672 assert_eq!(dirs.len(), 2);
673 let names: Vec<&str> = dirs.iter().map(|d| d.name.as_str()).collect();
674 assert!(names.contains(&".venv"));
675 assert!(names.contains(&"my_env"));
676 }
677
678 #[test]
679 fn test_bloat_dirs_absent() {
680 let dir = tempdir().unwrap();
681 let dirs = Venv.bloat_dirs(dir.path());
682 assert!(dirs.is_empty());
683 }
684
685 #[test]
686 fn test_bloat_dirs_ignores_non_venv_dirs() {
687 let dir = tempdir().unwrap();
688 fs::create_dir(dir.path().join("src")).unwrap();
690 make_venv(dir.path(), ".venv");
691 let dirs = Venv.bloat_dirs(dir.path());
692 assert_eq!(dirs.len(), 1);
693 assert_eq!(dirs[0].name, ".venv");
694 }
695
696 fn install_package(root: &Path, venv: &str, name: &str, requires: &[&str]) {
700 let dist_info = root
701 .join(venv)
702 .join("Lib")
703 .join("site-packages")
704 .join(format!("{name}-1.0.0.dist-info"));
705 fs::create_dir_all(&dist_info).unwrap();
706 let mut metadata = format!("Metadata-Version: 2.1\nName: {name}\nVersion: 1.0.0\n");
707 for dep in requires {
708 metadata.push_str(&format!("Requires-Dist: {dep}\n"));
709 }
710 fs::write(dist_info.join("METADATA"), metadata).unwrap();
711 }
712
713 #[test]
714 fn enforce_refuses_when_requirements_lists_nothing() {
715 let dir = tempdir().unwrap();
716 fs::write(dir.path().join("requirements.txt"), "# nothing pinned\n\n").unwrap();
717 make_venv(dir.path(), ".venv");
718
719 let err = Venv
720 .enforce_lockfile(dir.path(), EnforcePolicy::default())
721 .unwrap_err();
722 assert!(err.to_string().contains("lists no packages"));
723 }
724
725 #[test]
726 fn enforce_refuses_a_package_the_requirements_never_recorded() {
727 let dir = tempdir().unwrap();
730 fs::write(dir.path().join("requirements.txt"), "flask==3.0.0\n").unwrap();
731 make_venv(dir.path(), ".venv");
732 install_package(dir.path(), ".venv", "flask", &[]);
733 install_package(dir.path(), ".venv", "requests", &[]);
734
735 let err = Venv
736 .enforce_lockfile(dir.path(), EnforcePolicy::default())
737 .unwrap_err()
738 .to_string();
739 assert!(
740 err.contains("requests"),
741 "names the unrecorded package: {err}"
742 );
743 assert!(
744 !err.contains("flask"),
745 "must not blame the pinned one: {err}"
746 );
747 assert!(err.contains("pip freeze"), "says how to record it: {err}");
748 }
749
750 #[test]
751 fn enforce_accepts_transitive_dependencies_of_pinned_packages() {
752 let dir = tempdir().unwrap();
755 fs::write(dir.path().join("requirements.txt"), "requests==2.32.3\n").unwrap();
756 make_venv(dir.path(), ".venv");
757 install_package(
758 dir.path(),
759 ".venv",
760 "requests",
761 &["urllib3 (>=1.21.1)", "charset-normalizer"],
762 );
763 install_package(dir.path(), ".venv", "urllib3", &[]);
764 install_package(dir.path(), ".venv", "charset_normalizer", &[]);
766
767 assert!(
768 Venv.enforce_lockfile(dir.path(), EnforcePolicy::default())
769 .is_ok()
770 );
771 }
772
773 #[test]
774 fn enforce_skips_the_comparison_when_requirements_cannot_be_parsed() {
775 let dir = tempdir().unwrap();
778 fs::write(
779 dir.path().join("requirements.txt"),
780 "-e ./local-package\nflask==3.0.0\n",
781 )
782 .unwrap();
783 make_venv(dir.path(), ".venv");
784 install_package(dir.path(), ".venv", "left-behind", &[]);
785
786 assert!(
787 Venv.enforce_lockfile(dir.path(), EnforcePolicy::default())
788 .is_ok()
789 );
790 }
791
792 #[test]
793 fn drift_names_the_venv_and_the_unrecorded_packages() {
794 let dir = tempdir().unwrap();
795 fs::write(dir.path().join("requirements.txt"), "requests==2.32.3\n").unwrap();
796 make_venv(dir.path(), ".venv");
797 install_package(dir.path(), ".venv", "requests", &[]);
798 install_package(dir.path(), ".venv", "sneaky-pkg", &[]);
799
800 let reports = Venv.drift(dir.path());
801 assert_eq!(reports.len(), 1);
802 assert_eq!(reports[0].directory, ".venv");
803 assert_eq!(reports[0].unrecorded, vec!["sneaky-pkg"]);
804 assert_eq!(reports[0].record_command, "pip freeze > requirements.txt");
805 }
806
807 #[test]
810 fn drift_does_not_flag_transitive_dependencies_of_pinned_packages() {
811 let dir = tempdir().unwrap();
812 fs::write(dir.path().join("requirements.txt"), "requests==2.32.3\n").unwrap();
813 make_venv(dir.path(), ".venv");
814 install_package(dir.path(), ".venv", "requests", &["urllib3"]);
815 install_package(dir.path(), ".venv", "urllib3", &[]);
816
817 assert!(Venv.drift(dir.path()).is_empty());
818 }
819}