use super::{BloatDir, EnforcePolicy, PackageManager, dir_size, run_command_with_timeout};
use anyhow::{Result, anyhow};
use std::collections::{HashMap, HashSet};
use std::fs;
use std::path::{Path, PathBuf};
const PYVENV_CFG: &str = "pyvenv.cfg";
const FOREIGN_PYTHON_LOCKFILES: [&str; 3] = ["poetry.lock", "Pipfile.lock", "pdm.lock"];
pub(super) const BASELINE_DISTRIBUTIONS: [&str; 4] =
["pip", "setuptools", "wheel", "pkg-resources"];
pub struct Venv;
fn find_venv_dirs(path: &Path) -> Vec<std::path::PathBuf> {
let mut found = Vec::new();
let Ok(entries) = fs::read_dir(path) else {
return found;
};
for entry in entries.flatten() {
let entry_path = entry.path();
if entry_path.is_dir() && entry_path.join(PYVENV_CFG).exists() {
found.push(entry_path);
}
}
found
}
fn is_poetry_project(path: &Path) -> bool {
fs::read_to_string(path.join("pyproject.toml"))
.map(|c| {
c.lines()
.any(|l| l.trim_start().starts_with("[tool.poetry"))
})
.unwrap_or(false)
}
pub(super) fn normalize_package_name(name: &str) -> String {
let mut out = String::with_capacity(name.len());
let mut last_dash = false;
for c in name.chars() {
if c == '-' || c == '_' || c == '.' {
if !last_dash {
out.push('-');
}
last_dash = true;
} else {
out.push(c.to_ascii_lowercase());
last_dash = false;
}
}
out
}
fn requirement_name(spec: &str) -> Option<String> {
let name: String = spec
.chars()
.take_while(|c| c.is_ascii_alphanumeric() || matches!(c, '-' | '_' | '.'))
.collect();
if name.is_empty() || name.starts_with('.') || spec.contains("://") && !spec.contains(" @ ") {
return None;
}
Some(normalize_package_name(&name))
}
fn requirement_names(file: &Path, visited: &mut Vec<PathBuf>) -> Option<HashSet<String>> {
if visited.len() >= 8 || visited.iter().any(|p| p == file) {
return None;
}
visited.push(file.to_path_buf());
let content = fs::read_to_string(file).ok()?;
let dir = file.parent()?;
let mut names = HashSet::new();
for raw in content.lines() {
let mut line = raw.trim();
if let Some(idx) = line.find(" #") {
line = &line[..idx];
}
let line = line.trim_end_matches('\\').trim();
if line.is_empty() || line.starts_with('#') {
continue;
}
if let Some(included) = line
.strip_prefix("-r ")
.or_else(|| line.strip_prefix("--requirement "))
.or_else(|| line.strip_prefix("-c "))
.or_else(|| line.strip_prefix("--constraint "))
{
names.extend(requirement_names(&dir.join(included.trim()), visited)?);
continue;
}
if line.starts_with('-') {
if line.starts_with("-e") || line.starts_with("--editable") {
return None;
}
continue;
}
let spec = line.split(" @ ").next().unwrap_or(line).trim();
names.insert(requirement_name(spec)?);
}
Some(names)
}
pub(super) fn installed_distributions(venv: &Path) -> Option<HashMap<String, Vec<String>>> {
let mut site_packages: Vec<PathBuf> = Vec::new();
let windows_layout = venv.join("Lib").join("site-packages");
if windows_layout.is_dir() {
site_packages.push(windows_layout);
}
for lib in ["lib", "lib64"] {
let Ok(entries) = fs::read_dir(venv.join(lib)) else {
continue;
};
for entry in entries.flatten() {
let sp = entry.path().join("site-packages");
if sp.is_dir() {
site_packages.push(sp);
}
}
}
if site_packages.is_empty() {
return None;
}
let mut installed = HashMap::new();
for sp in site_packages {
let Ok(entries) = fs::read_dir(&sp) else {
continue;
};
for entry in entries.flatten() {
let file_name = entry.file_name().to_string_lossy().into_owned();
let Some(stem) = file_name
.strip_suffix(".dist-info")
.or_else(|| file_name.strip_suffix(".egg-info"))
else {
continue;
};
let name = stem.rsplit_once('-').map(|(n, _)| n).unwrap_or(stem);
installed.insert(
normalize_package_name(name),
declared_dependencies(&entry.path()),
);
}
}
Some(installed)
}
fn declared_dependencies(dist_info: &Path) -> Vec<String> {
let Ok(metadata) = fs::read_to_string(dist_info.join("METADATA")) else {
return Vec::new();
};
let mut deps = Vec::new();
for line in metadata.lines() {
if line.is_empty() {
break;
}
if let Some(spec) = line.strip_prefix("Requires-Dist:") {
if let Some(name) = requirement_name(spec.trim()) {
deps.push(name);
}
}
}
deps
}
fn venv_python_version(venv: &Path) -> Option<(u64, u64)> {
let cfg = fs::read_to_string(venv.join(PYVENV_CFG)).ok()?;
for line in cfg.lines() {
let Some((key, value)) = line.split_once('=') else {
continue;
};
if matches!(key.trim(), "version" | "version_info") {
let mut parts = value.trim().split('.');
return Some((parts.next()?.parse().ok()?, parts.next()?.parse().ok()?));
}
}
None
}
fn path_python_version() -> Option<(u64, u64)> {
let output = std::process::Command::new(super::resolve_program("python"))
.arg("--version")
.stdin(std::process::Stdio::null())
.output()
.ok()?;
if !output.status.success() {
return None;
}
let stdout = String::from_utf8_lossy(&output.stdout).trim().to_string();
let text = if stdout.is_empty() {
String::from_utf8_lossy(&output.stderr).trim().to_string()
} else {
stdout
};
let version = text.split_whitespace().nth(1)?;
let mut parts = version.split('.');
Some((parts.next()?.parse().ok()?, parts.next()?.parse().ok()?))
}
fn warn_about_restore_surprises(path: &Path, venvs: &[PathBuf]) {
if venvs.len() > 1 {
crate::output::print_warning(&format!(
"`{}` has {} virtual environments, all rebuilt from one requirements.txt. \
Each restores under its own recorded name; a plain `devp restore` with no \
record rebuilds only `.venv`.",
path.display(),
venvs.len()
));
} else if let Some(venv) = venvs.first() {
let name = venv.file_name().map(|n| n.to_string_lossy().into_owned());
if let Some(name) = name {
if name != ".venv" {
crate::output::print_info(&format!(
"The environment at `{}` is named `{name}` — `devp restore --last-run` \
recreates that name, but a restore with no record creates `.venv`.",
venv.display()
));
}
}
}
let on_path = path_python_version();
for venv in venvs {
if let (Some(built_with), Some(available)) = (venv_python_version(venv), on_path) {
if built_with != available {
crate::output::print_warning(&format!(
"`{}` was built with Python {}.{}, but `python` on PATH is {}.{} — a \
restore would rebuild it on that interpreter instead, and pinned \
wheels may not exist for it.",
venv.display(),
built_with.0,
built_with.1,
available.0,
available.1
));
}
}
}
}
fn unrecorded_packages(
installed: &HashMap<String, Vec<String>>,
pinned: &HashSet<String>,
) -> Vec<String> {
let mut reachable: HashSet<String> = HashSet::new();
let mut queue: Vec<String> = pinned.iter().cloned().collect();
queue.extend(BASELINE_DISTRIBUTIONS.iter().map(|s| (*s).to_string()));
while let Some(name) = queue.pop() {
if !reachable.insert(name.clone()) {
continue;
}
if let Some(deps) = installed.get(&name) {
queue.extend(deps.iter().cloned());
}
}
let mut extras: Vec<String> = installed
.keys()
.filter(|name| !reachable.contains(*name))
.cloned()
.collect();
extras.sort();
extras
}
impl PackageManager for Venv {
fn name(&self) -> &'static str {
"venv"
}
fn detect(&self, path: &Path) -> bool {
let req_txt = path.join("requirements.txt");
let uv_lock = path.join("uv.lock");
if !req_txt.exists() || uv_lock.exists() {
return false;
}
if FOREIGN_PYTHON_LOCKFILES
.iter()
.any(|f| path.join(f).exists())
|| is_poetry_project(path)
{
return false;
}
!find_venv_dirs(path).is_empty()
}
fn bloat_dirs(&self, path: &Path) -> Vec<BloatDir> {
find_venv_dirs(path)
.into_iter()
.map(|venv_path| {
let name = venv_path
.file_name()
.map(|n| n.to_string_lossy().to_string())
.unwrap_or_else(|| venv_path.display().to_string());
let size = dir_size(&venv_path);
BloatDir {
name,
path: venv_path,
size_bytes: size,
shared_bytes: 0,
}
})
.collect()
}
fn enforce_lockfile(&self, path: &Path, _policy: EnforcePolicy) -> Result<()> {
let req_txt = path.join("requirements.txt");
if !req_txt.exists() {
return Err(anyhow!("requirements.txt missing"));
}
let has_requirements = fs::read_to_string(&req_txt)
.map(|c| {
c.lines()
.any(|l| !l.trim().is_empty() && !l.trim_start().starts_with('#'))
})
.unwrap_or(false);
if !has_requirements {
return Err(anyhow!(
"requirements.txt at `{}` lists no packages — the virtual environment \
could not be rebuilt after deletion. Populate it with `pip freeze > requirements.txt`.",
req_txt.display()
));
}
let venvs = find_venv_dirs(path);
warn_about_restore_surprises(path, &venvs);
if let Some(pinned) = requirement_names(&req_txt, &mut Vec::new()) {
for venv in venvs {
let Some(installed) = installed_distributions(&venv) else {
continue;
};
let extras = unrecorded_packages(&installed, &pinned);
if extras.is_empty() {
continue;
}
let shown = extras
.iter()
.take(10)
.cloned()
.collect::<Vec<_>>()
.join(", ");
let suffix = if extras.len() > 10 {
format!(", … and {} more", extras.len() - 10)
} else {
String::new()
};
return Err(anyhow!(
"`{}` holds {} package(s) that requirements.txt does not account for \
({shown}{suffix}). Deleting the environment would lose them with no \
way back. Record them first: `pip freeze > requirements.txt`.",
venv.display(),
extras.len()
));
}
}
Ok(())
}
fn restore(&self, path: &Path, timeout: std::time::Duration) -> Result<()> {
self.restore_named(path, ".venv", timeout)
}
fn restore_named(
&self,
path: &Path,
dir_name: &str,
timeout: std::time::Duration,
) -> Result<()> {
let dir_name = if dir_name.is_empty()
|| dir_name == "."
|| dir_name == ".."
|| dir_name.contains(['/', '\\'])
{
".venv"
} else {
dir_name
};
run_command_with_timeout("python", &["-m", "venv", dir_name], path, timeout)?;
#[cfg(windows)]
let python = path.join(dir_name).join("Scripts").join("python.exe");
#[cfg(not(windows))]
let python = path.join(dir_name).join("bin").join("python");
run_command_with_timeout(
&python.to_string_lossy(),
&["-m", "pip", "install", "-r", "requirements.txt"],
path,
timeout,
)
}
fn lockfiles(&self) -> &'static [&'static str] {
&["requirements.txt"]
}
fn drift(&self, path: &Path) -> Vec<super::DriftReport> {
let Some(pinned) = requirement_names(&path.join("requirements.txt"), &mut Vec::new())
else {
return Vec::new();
};
let mut reports = Vec::new();
for venv in find_venv_dirs(path) {
let Some(installed) = installed_distributions(&venv) else {
continue;
};
let extras = unrecorded_packages(&installed, &pinned);
if extras.is_empty() {
continue;
}
reports.push(super::DriftReport {
directory: venv
.file_name()
.map(|n| n.to_string_lossy().into_owned())
.unwrap_or_else(|| venv.display().to_string()),
unrecorded: extras,
record_command: "pip freeze > requirements.txt",
});
}
reports
}
}
#[cfg(test)]
mod tests {
use super::*;
use std::fs::{self, File};
use tempfile::tempdir;
fn make_venv(dir: &Path, name: &str) {
let venv = dir.join(name);
fs::create_dir(&venv).unwrap();
File::create(venv.join(PYVENV_CFG)).unwrap();
}
#[test]
fn test_name() {
assert_eq!(Venv.name(), "venv");
}
#[test]
fn test_detect_positive_dot_venv() {
let dir = tempdir().unwrap();
File::create(dir.path().join("requirements.txt")).unwrap();
make_venv(dir.path(), ".venv");
assert!(Venv.detect(dir.path()));
}
#[test]
fn test_detect_positive_venv() {
let dir = tempdir().unwrap();
File::create(dir.path().join("requirements.txt")).unwrap();
make_venv(dir.path(), "venv");
assert!(Venv.detect(dir.path()));
}
#[test]
fn test_detect_positive_custom_name() {
let dir = tempdir().unwrap();
File::create(dir.path().join("requirements.txt")).unwrap();
make_venv(dir.path(), "my_env");
assert!(Venv.detect(dir.path()));
}
#[test]
fn test_detect_positive_env() {
let dir = tempdir().unwrap();
File::create(dir.path().join("requirements.txt")).unwrap();
make_venv(dir.path(), "env");
assert!(Venv.detect(dir.path()));
}
#[test]
fn test_detect_negative_no_req() {
let dir = tempdir().unwrap();
make_venv(dir.path(), ".venv");
assert!(!Venv.detect(dir.path()));
}
#[test]
fn test_detect_negative_no_env() {
let dir = tempdir().unwrap();
File::create(dir.path().join("requirements.txt")).unwrap();
fs::create_dir(dir.path().join("not_a_venv")).unwrap();
assert!(!Venv.detect(dir.path()));
}
#[test]
fn test_detect_negative_uv_lock() {
let dir = tempdir().unwrap();
File::create(dir.path().join("requirements.txt")).unwrap();
File::create(dir.path().join("uv.lock")).unwrap();
make_venv(dir.path(), ".venv");
assert!(!Venv.detect(dir.path()));
}
#[test]
fn test_bloat_dirs_present() {
let dir = tempdir().unwrap();
make_venv(dir.path(), ".venv");
make_venv(dir.path(), "my_env");
let dirs = Venv.bloat_dirs(dir.path());
assert_eq!(dirs.len(), 2);
let names: Vec<&str> = dirs.iter().map(|d| d.name.as_str()).collect();
assert!(names.contains(&".venv"));
assert!(names.contains(&"my_env"));
}
#[test]
fn test_bloat_dirs_absent() {
let dir = tempdir().unwrap();
let dirs = Venv.bloat_dirs(dir.path());
assert!(dirs.is_empty());
}
#[test]
fn test_bloat_dirs_ignores_non_venv_dirs() {
let dir = tempdir().unwrap();
fs::create_dir(dir.path().join("src")).unwrap();
make_venv(dir.path(), ".venv");
let dirs = Venv.bloat_dirs(dir.path());
assert_eq!(dirs.len(), 1);
assert_eq!(dirs[0].name, ".venv");
}
fn install_package(root: &Path, venv: &str, name: &str, requires: &[&str]) {
let dist_info = root
.join(venv)
.join("Lib")
.join("site-packages")
.join(format!("{name}-1.0.0.dist-info"));
fs::create_dir_all(&dist_info).unwrap();
let mut metadata = format!("Metadata-Version: 2.1\nName: {name}\nVersion: 1.0.0\n");
for dep in requires {
metadata.push_str(&format!("Requires-Dist: {dep}\n"));
}
fs::write(dist_info.join("METADATA"), metadata).unwrap();
}
#[test]
fn enforce_refuses_when_requirements_lists_nothing() {
let dir = tempdir().unwrap();
fs::write(dir.path().join("requirements.txt"), "# nothing pinned\n\n").unwrap();
make_venv(dir.path(), ".venv");
let err = Venv
.enforce_lockfile(dir.path(), EnforcePolicy::default())
.unwrap_err();
assert!(err.to_string().contains("lists no packages"));
}
#[test]
fn enforce_refuses_a_package_the_requirements_never_recorded() {
let dir = tempdir().unwrap();
fs::write(dir.path().join("requirements.txt"), "flask==3.0.0\n").unwrap();
make_venv(dir.path(), ".venv");
install_package(dir.path(), ".venv", "flask", &[]);
install_package(dir.path(), ".venv", "requests", &[]);
let err = Venv
.enforce_lockfile(dir.path(), EnforcePolicy::default())
.unwrap_err()
.to_string();
assert!(
err.contains("requests"),
"names the unrecorded package: {err}"
);
assert!(
!err.contains("flask"),
"must not blame the pinned one: {err}"
);
assert!(err.contains("pip freeze"), "says how to record it: {err}");
}
#[test]
fn enforce_accepts_transitive_dependencies_of_pinned_packages() {
let dir = tempdir().unwrap();
fs::write(dir.path().join("requirements.txt"), "requests==2.32.3\n").unwrap();
make_venv(dir.path(), ".venv");
install_package(
dir.path(),
".venv",
"requests",
&["urllib3 (>=1.21.1)", "charset-normalizer"],
);
install_package(dir.path(), ".venv", "urllib3", &[]);
install_package(dir.path(), ".venv", "charset_normalizer", &[]);
assert!(
Venv.enforce_lockfile(dir.path(), EnforcePolicy::default())
.is_ok()
);
}
#[test]
fn enforce_skips_the_comparison_when_requirements_cannot_be_parsed() {
let dir = tempdir().unwrap();
fs::write(
dir.path().join("requirements.txt"),
"-e ./local-package\nflask==3.0.0\n",
)
.unwrap();
make_venv(dir.path(), ".venv");
install_package(dir.path(), ".venv", "left-behind", &[]);
assert!(
Venv.enforce_lockfile(dir.path(), EnforcePolicy::default())
.is_ok()
);
}
#[test]
fn drift_names_the_venv_and_the_unrecorded_packages() {
let dir = tempdir().unwrap();
fs::write(dir.path().join("requirements.txt"), "requests==2.32.3\n").unwrap();
make_venv(dir.path(), ".venv");
install_package(dir.path(), ".venv", "requests", &[]);
install_package(dir.path(), ".venv", "sneaky-pkg", &[]);
let reports = Venv.drift(dir.path());
assert_eq!(reports.len(), 1);
assert_eq!(reports[0].directory, ".venv");
assert_eq!(reports[0].unrecorded, vec!["sneaky-pkg"]);
assert_eq!(reports[0].record_command, "pip freeze > requirements.txt");
}
#[test]
fn drift_does_not_flag_transitive_dependencies_of_pinned_packages() {
let dir = tempdir().unwrap();
fs::write(dir.path().join("requirements.txt"), "requests==2.32.3\n").unwrap();
make_venv(dir.path(), ".venv");
install_package(dir.path(), ".venv", "requests", &["urllib3"]);
install_package(dir.path(), ".venv", "urllib3", &[]);
assert!(Venv.drift(dir.path()).is_empty());
}
}