Skip to content

Commit

Permalink
Preserve seed packages for non-Puffin-created virtualenvs (#535)
Browse files Browse the repository at this point in the history
## Summary

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

Closes #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.
  • Loading branch information
charliermarsh authored Dec 4, 2023
1 parent 77b3921 commit 95b8316
Show file tree
Hide file tree
Showing 5 changed files with 94 additions and 3 deletions.
16 changes: 13 additions & 3 deletions crates/puffin-installer/src/plan.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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 {
Expand Down
10 changes: 10 additions & 0 deletions crates/puffin-installer/src/site_packages.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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 {
Expand Down
60 changes: 60 additions & 0 deletions crates/puffin-interpreter/src/cfg.rs
Original file line number Diff line number Diff line change
@@ -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),
}
3 changes: 3 additions & 0 deletions crates/puffin-interpreter/src/lib.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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;
Expand Down Expand Up @@ -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),
}
8 changes: 8 additions & 0 deletions crates/puffin-interpreter/src/virtual_env.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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};

Expand Down Expand Up @@ -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
Expand Down

0 comments on commit 95b8316

Please sign in to comment.