From b90140e1bc5d2aea08a2699126defa32969ca456 Mon Sep 17 00:00:00 2001 From: Charlie Marsh Date: Mon, 9 Oct 2023 14:14:33 -0400 Subject: [PATCH] Add support for wheel uninstalls (#77) Closes #36. --- crates/install-wheel-rs/src/lib.rs | 14 +- crates/install-wheel-rs/src/uninstall.rs | 129 ++++++++++++++++++ crates/puffin-cli/src/commands/freeze.rs | 4 +- crates/puffin-cli/src/commands/mod.rs | 2 + crates/puffin-cli/src/commands/sync.rs | 8 +- crates/puffin-cli/src/commands/uninstall.rs | 25 ++++ crates/puffin-cli/src/main.rs | 21 +++ crates/puffin-installer/src/install.rs | 10 +- crates/puffin-installer/src/lib.rs | 2 + crates/puffin-installer/src/uninstall.rs | 36 +++++ .../puffin-interpreter/src/site_packages.rs | 32 ++++- 11 files changed, 256 insertions(+), 27 deletions(-) create mode 100644 crates/install-wheel-rs/src/uninstall.rs create mode 100644 crates/puffin-cli/src/commands/uninstall.rs create mode 100644 crates/puffin-installer/src/uninstall.rs diff --git a/crates/install-wheel-rs/src/lib.rs b/crates/install-wheel-rs/src/lib.rs index 65428e3b6862..6c848a0f5be9 100644 --- a/crates/install-wheel-rs/src/lib.rs +++ b/crates/install-wheel-rs/src/lib.rs @@ -1,17 +1,16 @@ //! Takes a wheel and installs it into a venv.. use std::io; -use std::io::{Read, Seek}; use platform_info::PlatformInfoError; use thiserror::Error; use zip::result::ZipError; -use zip::ZipArchive; pub use install_location::{normalize_name, InstallLocation, LockedDir}; use platform_host::{Arch, Os}; pub use record::RecordEntry; pub use script::Script; +pub use uninstall::uninstall_wheel; pub use wheel::{ get_script_launcher, install_wheel, parse_key_value_file, read_record_file, relative_to, SHEBANG_PYTHON, @@ -24,6 +23,7 @@ mod record; #[cfg(any(target_os = "macos", target_os = "ios"))] mod reflink; mod script; +mod uninstall; pub mod unpacked; mod wheel; @@ -74,13 +74,3 @@ impl Error { } } } - -pub fn do_thing(reader: impl Read + Seek) -> Result<(), Error> { - let x = tempfile::tempdir()?; - let mut archive = - ZipArchive::new(reader).map_err(|err| Error::from_zip_error("(index)".to_string(), err))?; - - archive.extract(x.path()).unwrap(); - - Ok(()) -} diff --git a/crates/install-wheel-rs/src/uninstall.rs b/crates/install-wheel-rs/src/uninstall.rs new file mode 100644 index 000000000000..4f716f730700 --- /dev/null +++ b/crates/install-wheel-rs/src/uninstall.rs @@ -0,0 +1,129 @@ +use std::collections::BTreeSet; +use std::path::{Component, Path, PathBuf}; + +use fs_err as fs; +use fs_err::File; +use tracing::debug; + +use crate::{read_record_file, Error}; + +/// Uninstall the wheel represented by the given `dist_info` directory. +pub fn uninstall_wheel(dist_info: &Path) -> Result { + let Some(site_packages) = dist_info.parent() else { + return Err(Error::BrokenVenv( + "dist-info directory is not in a site-packages directory".to_string(), + )); + }; + + // Read the RECORD file. + let mut record_file = File::open(dist_info.join("RECORD"))?; + let record = read_record_file(&mut record_file)?; + + let mut file_count = 0usize; + let mut dir_count = 0usize; + + // Uninstall the files, keeping track of any directories that are left empty. + let mut visited = BTreeSet::new(); + for entry in &record { + let path = site_packages.join(&entry.path); + match fs::remove_file(&path) { + Ok(()) => { + debug!("Removed file: {}", path.display()); + file_count += 1; + } + Err(err) if err.kind() == std::io::ErrorKind::NotFound => {} + Err(err) => return Err(err.into()), + } + + if let Some(parent) = path.parent() { + visited.insert(normalize_path(parent)); + } + } + + // If any directories were left empty, remove them. Iterate in reverse order such that we visit + // the deepest directories first. + for path in visited.iter().rev() { + // No need to look at directories outside of `site-packages` (like `bin`). + if !path.starts_with(site_packages) { + continue; + } + + // Iterate up the directory tree, removing any empty directories. It's insufficient to + // rely on `visited` alone here, because we may end up removing a directory whose parent + // directory doesn't contain any files, leaving the _parent_ directory empty. + let mut path = path.as_path(); + loop { + // If we reach the site-packages directory, we're done. + if path == site_packages { + break; + } + + // Try to read from the directory. If it doesn't exist, assume we deleted it in a + // previous iteration. + let mut read_dir = match fs::read_dir(path) { + Ok(read_dir) => read_dir, + Err(err) if err.kind() == std::io::ErrorKind::NotFound => break, + Err(err) => return Err(err.into()), + }; + + // If the directory is not empty, we're done. + if read_dir.next().is_some() { + break; + } + + fs::remove_dir(path)?; + + debug!("Removed directory: {}", path.display()); + dir_count += 1; + + if let Some(parent) = path.parent() { + path = parent; + } else { + break; + } + } + } + + Ok(Uninstall { + file_count, + dir_count, + }) +} + +#[derive(Debug)] +pub struct Uninstall { + /// The number of files that were removed during the uninstallation. + pub file_count: usize, + /// The number of directories that were removed during the uninstallation. + pub dir_count: usize, +} + +/// Normalize a path, removing things like `.` and `..`. +/// +/// Source: +fn normalize_path(path: &Path) -> PathBuf { + let mut components = path.components().peekable(); + let mut ret = if let Some(c @ Component::Prefix(..)) = components.peek().copied() { + components.next(); + PathBuf::from(c.as_os_str()) + } else { + PathBuf::new() + }; + + for component in components { + match component { + Component::Prefix(..) => unreachable!(), + Component::RootDir => { + ret.push(component.as_os_str()); + } + Component::CurDir => {} + Component::ParentDir => { + ret.pop(); + } + Component::Normal(c) => { + ret.push(c); + } + } + } + ret +} diff --git a/crates/puffin-cli/src/commands/freeze.rs b/crates/puffin-cli/src/commands/freeze.rs index 66f41a94062d..ce09bae92069 100644 --- a/crates/puffin-cli/src/commands/freeze.rs +++ b/crates/puffin-cli/src/commands/freeze.rs @@ -20,10 +20,10 @@ pub(crate) async fn freeze(cache: Option<&Path>) -> Result { // Build the installed index. let site_packages = SitePackages::from_executable(&python).await?; - for (name, version) in site_packages.iter() { + for (name, dist_info) in site_packages.iter() { #[allow(clippy::print_stdout)] { - println!("{name}=={version}"); + println!("{}=={}", name, dist_info.version()); } } diff --git a/crates/puffin-cli/src/commands/mod.rs b/crates/puffin-cli/src/commands/mod.rs index 894c2796ef38..7fa4ffa89ecf 100644 --- a/crates/puffin-cli/src/commands/mod.rs +++ b/crates/puffin-cli/src/commands/mod.rs @@ -5,11 +5,13 @@ pub(crate) use clean::clean; pub(crate) use compile::compile; pub(crate) use freeze::freeze; pub(crate) use sync::{sync, SyncFlags}; +pub(crate) use uninstall::uninstall; mod clean; mod compile; mod freeze; mod sync; +mod uninstall; #[derive(Copy, Clone)] pub(crate) enum ExitStatus { diff --git a/crates/puffin-cli/src/commands/sync.rs b/crates/puffin-cli/src/commands/sync.rs index 3c9ffa382108..4f8e55e54052 100644 --- a/crates/puffin-cli/src/commands/sync.rs +++ b/crates/puffin-cli/src/commands/sync.rs @@ -66,8 +66,12 @@ pub(crate) async fn sync(src: &Path, cache: Option<&Path>, flags: SyncFlags) -> let package = PackageName::normalize(&requirement.name); // Filter out already-installed packages. - if let Some(version) = site_packages.get(&package) { - info!("Requirement already satisfied: {package} ({version})"); + if let Some(dist_info) = site_packages.get(&package) { + info!( + "Requirement already satisfied: {} ({})", + package, + dist_info.version() + ); return None; } diff --git a/crates/puffin-cli/src/commands/uninstall.rs b/crates/puffin-cli/src/commands/uninstall.rs new file mode 100644 index 000000000000..e74283edf4f5 --- /dev/null +++ b/crates/puffin-cli/src/commands/uninstall.rs @@ -0,0 +1,25 @@ +use std::path::Path; + +use anyhow::Result; +use tracing::debug; + +use platform_host::Platform; +use puffin_interpreter::PythonExecutable; + +use crate::commands::ExitStatus; + +/// Uninstall a package from the current environment. +pub(crate) async fn uninstall(name: &str, cache: Option<&Path>) -> Result { + // Detect the current Python interpreter. + let platform = Platform::current()?; + let python = PythonExecutable::from_env(platform, cache)?; + debug!( + "Using Python interpreter: {}", + python.executable().display() + ); + + // Uninstall the package from the current environment. + puffin_installer::uninstall(name, &python).await?; + + Ok(ExitStatus::Success) +} diff --git a/crates/puffin-cli/src/main.rs b/crates/puffin-cli/src/main.rs index 222d9ce868bf..2c48d41603a2 100644 --- a/crates/puffin-cli/src/main.rs +++ b/crates/puffin-cli/src/main.rs @@ -36,6 +36,8 @@ enum Commands { Clean, /// Enumerate the installed packages in the current environment. Freeze(FreezeArgs), + /// Uninstall a package. + Uninstall(UninstallArgs), } #[derive(Args)] @@ -69,6 +71,16 @@ struct FreezeArgs { no_cache: bool, } +#[derive(Args)] +struct UninstallArgs { + /// The name of the package to uninstall. + name: String, + + /// Avoid reading from or writing to the cache. + #[arg(long)] + no_cache: bool, +} + #[tokio::main] async fn main() -> ExitCode { let cli = Cli::parse(); @@ -116,6 +128,15 @@ async fn main() -> ExitCode { ) .await } + Commands::Uninstall(args) => { + commands::uninstall( + &args.name, + dirs.as_ref() + .map(ProjectDirs::cache_dir) + .filter(|_| !args.no_cache), + ) + .await + } }; match result { diff --git a/crates/puffin-installer/src/install.rs b/crates/puffin-installer/src/install.rs index 74e90f57ddde..742301d9e7da 100644 --- a/crates/puffin-installer/src/install.rs +++ b/crates/puffin-installer/src/install.rs @@ -10,7 +10,6 @@ use tracing::{debug, info}; use url::Url; use zip::ZipArchive; -use install_wheel_rs::{unpacked, InstallLocation}; use puffin_client::PypiClient; use puffin_interpreter::PythonExecutable; @@ -101,7 +100,10 @@ pub async fn install( ); // Phase 3: Install each wheel. - let location = InstallLocation::new(python.venv().to_path_buf(), python.simple_version()); + let location = install_wheel_rs::InstallLocation::new( + python.venv().to_path_buf(), + python.simple_version(), + ); let locked_dir = location.acquire_lock()?; for wheel in wheels { @@ -112,10 +114,10 @@ pub async fn install( || staging.path().join(&id), |wheel_cache| wheel_cache.entry(&id), ); - unpacked::install_wheel(&locked_dir, &dir)?; + install_wheel_rs::unpacked::install_wheel(&locked_dir, &dir)?; } Distribution::Local(local) => { - unpacked::install_wheel(&locked_dir, local.path())?; + install_wheel_rs::unpacked::install_wheel(&locked_dir, local.path())?; } } } diff --git a/crates/puffin-installer/src/lib.rs b/crates/puffin-installer/src/lib.rs index 21532e843ac3..c793bf765354 100644 --- a/crates/puffin-installer/src/lib.rs +++ b/crates/puffin-installer/src/lib.rs @@ -1,9 +1,11 @@ pub use distribution::{Distribution, LocalDistribution, RemoteDistribution}; pub use index::LocalIndex; pub use install::install; +pub use uninstall::uninstall; mod cache; mod distribution; mod index; mod install; +mod uninstall; mod vendor; diff --git a/crates/puffin-installer/src/uninstall.rs b/crates/puffin-installer/src/uninstall.rs new file mode 100644 index 000000000000..9c9624977bbd --- /dev/null +++ b/crates/puffin-installer/src/uninstall.rs @@ -0,0 +1,36 @@ +use anyhow::{anyhow, Result}; +use tracing::info; + +use puffin_interpreter::PythonExecutable; +use puffin_package::package_name::PackageName; + +/// Uninstall a package from the specified Python environment. +pub async fn uninstall(name: &str, python: &PythonExecutable) -> Result<()> { + // Index the current `site-packages` directory. + let site_packages = puffin_interpreter::SitePackages::from_executable(python).await?; + + // Locate the package in the environment. + let Some(dist_info) = site_packages.get(&PackageName::normalize(name)) else { + return Err(anyhow!("Package not installed: {}", name)); + }; + + // Uninstall the package from the environment. + let uninstall = tokio::task::spawn_blocking({ + let path = dist_info.path().to_owned(); + move || install_wheel_rs::uninstall_wheel(&path) + }) + .await??; + + // Print a summary of the uninstallation. + match (uninstall.file_count, uninstall.dir_count) { + (0, 0) => info!("No files found"), + (1, 0) => info!("Removed 1 file"), + (0, 1) => info!("Removed 1 directory"), + (1, 1) => info!("Removed 1 file and 1 directory"), + (file_count, 0) => info!("Removed {file_count} files"), + (0, dir_count) => info!("Removed {dir_count} directories"), + (file_count, dir_count) => info!("Removed {file_count} files and {dir_count} directories"), + } + + Ok(()) +} diff --git a/crates/puffin-interpreter/src/site_packages.rs b/crates/puffin-interpreter/src/site_packages.rs index cbe302e078bb..6569cf72c733 100644 --- a/crates/puffin-interpreter/src/site_packages.rs +++ b/crates/puffin-interpreter/src/site_packages.rs @@ -1,5 +1,5 @@ use std::collections::BTreeMap; -use std::path::Path; +use std::path::{Path, PathBuf}; use std::str::FromStr; use anyhow::{anyhow, Result}; @@ -10,7 +10,7 @@ use puffin_package::package_name::PackageName; use crate::PythonExecutable; #[derive(Debug, Default)] -pub struct SitePackages(BTreeMap); +pub struct SitePackages(BTreeMap); impl SitePackages { /// Build an index of installed packages from the given Python executable. @@ -21,7 +21,7 @@ impl SitePackages { while let Some(entry) = dir.next_entry().await? { if entry.file_type().await?.is_dir() { if let Some(dist_info) = DistInfo::try_from_path(&entry.path())? { - index.insert(dist_info.name, dist_info.version); + index.insert(dist_info.name().clone(), dist_info); } } } @@ -30,20 +30,21 @@ impl SitePackages { } /// Returns an iterator over the installed packages. - pub fn iter(&self) -> impl Iterator { + pub fn iter(&self) -> impl Iterator { self.0.iter() } /// Returns the version of the given package, if it is installed. - pub fn get(&self, name: &PackageName) -> Option<&Version> { + pub fn get(&self, name: &PackageName) -> Option<&DistInfo> { self.0.get(name) } } #[derive(Debug)] -struct DistInfo { +pub struct DistInfo { name: PackageName, version: Version, + path: PathBuf, } impl DistInfo { @@ -64,10 +65,27 @@ impl DistInfo { let name = PackageName::normalize(name); let version = Version::from_str(version).map_err(|err| anyhow!(err))?; + let path = path.to_path_buf(); - return Ok(Some(DistInfo { name, version })); + return Ok(Some(DistInfo { + name, + version, + path, + })); } Ok(None) } + + pub fn name(&self) -> &PackageName { + &self.name + } + + pub fn version(&self) -> &Version { + &self.version + } + + pub fn path(&self) -> &Path { + &self.path + } }