Preserve seed packages for non-Puffin-created virtualenvs (#535)

## Summary

This PR modifies the install plan to avoid removing seed packages if the
virtual environment was created by anyone other than Puffin.

Closes https://github.com/astral-sh/puffin/issues/414.

## Test Plan

- Ran: `virtualenv .venv`.
- Ran: `cargo run -p puffin-cli -- pip-sync
scripts/benchmarks/requirements.txt --verbose --no-cache`.
- Verified that `pip` et al were not removed, and that the logging
including a message around preserving seed packages.
This commit is contained in:
Charlie Marsh 2023-12-04 09:31:00 -05:00 committed by GitHub
parent 77b3921b7a
commit 95b8316023
No known key found for this signature in database
GPG key ID: 4AEE18F83AFDEB23
5 changed files with 94 additions and 3 deletions

View file

@ -142,9 +142,19 @@ impl InstallPlan {
}
// Remove any unnecessary packages.
for (_package, dist_info) in site_packages {
debug!("Unnecessary package: {dist_info}");
extraneous.push(dist_info);
if !site_packages.is_empty() {
// If Puffin created the virtual environment, then remove all packages, regardless of
// whether they're considered "seed" packages.
let seed_packages = !venv.cfg().is_ok_and(|cfg| cfg.is_gourgeist());
for (package, dist_info) in site_packages {
if seed_packages && matches!(package.as_ref(), "pip" | "setuptools" | "wheel") {
debug!("Preserving seed package: {dist_info}");
continue;
}
debug!("Unnecessary package: {dist_info}");
extraneous.push(dist_info);
}
}
Ok(InstallPlan {

View file

@ -45,6 +45,16 @@ impl SitePackages {
pub fn remove(&mut self, name: &PackageName) -> Option<InstalledDist> {
self.0.remove(name)
}
/// Returns `true` if there are no installed packages.
pub fn is_empty(&self) -> bool {
self.0.is_empty()
}
/// Returns the number of installed packages.
pub fn len(&self) -> usize {
self.0.len()
}
}
impl IntoIterator for SitePackages {

View file

@ -0,0 +1,60 @@
use std::path::Path;
use fs_err as fs;
use thiserror::Error;
#[derive(Debug, Clone)]
pub struct Configuration {
/// The version of the `virtualenv` package used to create the virtual environment, if any.
pub(crate) virtualenv: bool,
/// The version of the `gourgeist` package used to create the virtual environment, if any.
pub(crate) gourgeist: bool,
}
impl Configuration {
/// Parse a `pyvenv.cfg` file into a [`Configuration`].
pub fn parse(cfg: impl AsRef<Path>) -> Result<Self, Error> {
let mut virtualenv = false;
let mut gourgeist = false;
// Per https://snarky.ca/how-virtual-environments-work/, the `pyvenv.cfg` file is not a
// valid INI file, and is instead expected to be parsed by partitioning each line on the
// first equals sign.
let content = fs::read_to_string(&cfg)?;
for line in content.lines() {
let Some((key, _value)) = line.split_once('=') else {
continue;
};
match key.trim() {
"virtualenv" => {
virtualenv = true;
}
"gourgeist" => {
gourgeist = true;
}
_ => {}
}
}
Ok(Self {
virtualenv,
gourgeist,
})
}
/// Returns true if the virtual environment was created with the `virtualenv` package.
pub fn is_virtualenv(&self) -> bool {
self.virtualenv
}
/// Returns true if the virtual environment was created with the `gourgeist` package.
pub fn is_gourgeist(&self) -> bool {
self.gourgeist
}
}
#[derive(Debug, Error)]
pub enum Error {
#[error(transparent)]
Io(#[from] std::io::Error),
}

View file

@ -7,6 +7,7 @@ use thiserror::Error;
pub use crate::interpreter::Interpreter;
pub use crate::virtual_env::Virtualenv;
mod cfg;
mod interpreter;
mod python_platform;
mod virtual_env;
@ -39,4 +40,6 @@ pub enum Error {
},
#[error("Failed to write to cache")]
Serde(#[from] serde_json::Error),
#[error("Failed to parse pyvenv.cfg")]
Cfg(#[from] cfg::Error),
}

View file

@ -6,6 +6,7 @@ use tracing::debug;
use platform_host::Platform;
use puffin_cache::Cache;
use crate::cfg::Configuration;
use crate::python_platform::PythonPlatform;
use crate::{Error, Interpreter};
@ -66,10 +67,17 @@ impl Virtualenv {
&self.root
}
/// Return the [`Interpreter`] for this virtual environment.
pub fn interpreter(&self) -> &Interpreter {
&self.interpreter
}
/// Return the [`Configuration`] for this virtual environment, as extracted from the
/// `pyvenv.cfg` file.
pub fn cfg(&self) -> Result<Configuration, Error> {
Ok(Configuration::parse(self.root.join("pyvenv.cfg"))?)
}
/// Returns the path to the `site-packages` directory inside a virtual environment.
pub fn site_packages(&self) -> PathBuf {
self.interpreter