Add support for wheel uninstalls (#77)

Closes #36.

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(())
-}

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<Uninstall, Error> {
+    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: <https://github.com/rust-lang/cargo/blob/b48c41aedbd69ee3990d62a0e2006edbb506a480/crates/cargo-util/src/paths.rs#L76C1-L109C2>
+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
+}

crates/puffin-cli/src/commands/freeze.rs

@@ -20,10 +20,10 @@ pub(crate) async fn freeze(cache: Option<&Path>) -> Result<ExitStatus> {
 
     // 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());
         }
     }
 

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 {

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;
             }
 

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<ExitStatus> {
+    // 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)
+}

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 {

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())?;
             }
         }
     }

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;

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(())
+}

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<PackageName, Version>);
+pub struct SitePackages(BTreeMap<PackageName, DistInfo>);
 
 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<Item = (&PackageName, &Version)> {
+    pub fn iter(&self) -> impl Iterator<Item = (&PackageName, &DistInfo)> {
         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
+    }
 }