From 6c8ce1d01394b64779ddbb547b704bf854bea099 Mon Sep 17 00:00:00 2001
From: Ahmed Ilyas <ahmed.n.ilyas@gmail.com>
Date: Sat, 6 Jul 2024 21:37:15 +0200
Subject: [PATCH] `uv cache prune` removes all cached environments (#4845)
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

## Summary

Resolves #4802

## Test Plan
- `cargo test`
```sh
❯ cargo run -- cache prune -v
Pruning cache at: /Users/ahmedilyas/Library/Caches/uv
No unused entries found
❯ cargo run -- tool run cowsay
warning: `uv tool run` is experimental and may change without warning.
Resolved 1 package in 182ms
Installed 1 package in 20ms
 + cowsay==6.1
usage: Cowsay [-h] [-c CHARACTER] -t TEXT [-v]
Cowsay: error: the following arguments are required: -t/--text
❯ cargo run -- cache prune -v
Pruning cache at: /Users/ahmedilyas/Library/Caches/uv
    0.793440s DEBUG uv_cache Removing dangling cache entry: /Users/ahmedilyas/Library/Caches/uv/environments-v1/095cd7c4c298a0d8
Removed 41 files (143.5KiB)
```
---
 crates/uv-cache/src/lib.rs     | 33 +++++++++++++++++----
 crates/uv/tests/cache_prune.rs | 53 ++++++++++++++++++++++++++++++++++
 crates/uv/tests/common/mod.rs  |  4 +++
 3 files changed, 84 insertions(+), 6 deletions(-)

diff --git a/crates/uv-cache/src/lib.rs b/crates/uv-cache/src/lib.rs
index a94dcbf60..f08100ea2 100644
--- a/crates/uv-cache/src/lib.rs
+++ b/crates/uv-cache/src/lib.rs
@@ -388,13 +388,34 @@ impl Cache {
             }
         }
 
-        for entry in fs::read_dir(self.bucket(CacheBucket::Archive))? {
-            let entry = entry?;
-            let path = entry.path().canonicalize()?;
-            if !references.contains(&path) {
-                debug!("Removing dangling cache entry: {}", path.display());
-                summary += rm_rf(path)?;
+        match fs::read_dir(self.bucket(CacheBucket::Archive)) {
+            Ok(entries) => {
+                for entry in entries {
+                    let entry = entry?;
+                    let path = entry.path().canonicalize()?;
+                    if !references.contains(&path) {
+                        debug!("Removing dangling cache entry: {}", path.display());
+                        summary += rm_rf(path)?;
+                    }
+                }
             }
+            Err(err) if err.kind() == io::ErrorKind::NotFound => (),
+            Err(err) => return Err(err),
+        }
+
+        // Third, remove any cached environments. These are never referenced by symlinks, so we can
+        // remove them directly.
+        match fs::read_dir(self.bucket(CacheBucket::Environments)) {
+            Ok(entries) => {
+                for entry in entries {
+                    let entry = entry?;
+                    let path = entry.path().canonicalize()?;
+                    debug!("Removing dangling cache entry: {}", path.display());
+                    summary += rm_rf(path)?;
+                }
+            }
+            Err(err) if err.kind() == io::ErrorKind::NotFound => (),
+            Err(err) => return Err(err),
         }
 
         Ok(summary)
diff --git a/crates/uv/tests/cache_prune.rs b/crates/uv/tests/cache_prune.rs
index fcab2202d..0da22defd 100644
--- a/crates/uv/tests/cache_prune.rs
+++ b/crates/uv/tests/cache_prune.rs
@@ -83,6 +83,59 @@ fn prune_stale_directory() -> Result<()> {
     Ok(())
 }
 
+/// `cache prune` should remove all cached environments from the cache.
+#[test]
+fn prune_cached_env() {
+    let context = TestContext::new("3.12").with_filtered_counts();
+    let tool_dir = context.temp_dir.child("tools");
+    let bin_dir = context.temp_dir.child("bin");
+
+    uv_snapshot!(context.filters(), context.tool_run()
+        .arg("pytest@8.0.0")
+        .arg("--version")
+        .env("UV_TOOL_DIR", tool_dir.as_os_str())
+        .env("XDG_BIN_HOME", bin_dir.as_os_str()), @r###"
+    success: true
+    exit_code: 0
+    ----- stdout -----
+    pytest 8.0.0
+
+    ----- stderr -----
+    warning: `uv tool run` is experimental and may change without warning.
+    Resolved [N] packages in [TIME]
+    Prepared [N] packages in [TIME]
+    Installed [N] packages in [TIME]
+     + iniconfig==2.0.0
+     + packaging==24.0
+     + pluggy==1.4.0
+     + pytest==8.0.0
+    "###);
+
+    let filters: Vec<_> = context
+        .filters()
+        .into_iter()
+        .chain([
+            // The cache entry does not have a stable key, so we filter it out
+            (
+                r"\[CACHE_DIR\](\\|\/)(.+)(\\|\/).*",
+                "[CACHE_DIR]/$2/[ENTRY]",
+            ),
+        ])
+        .collect();
+
+    uv_snapshot!(filters, prune_command(&context).arg("--verbose"), @r###"
+    success: true
+    exit_code: 0
+    ----- stdout -----
+
+    ----- stderr -----
+    DEBUG uv [VERSION] ([COMMIT] DATE)
+    Pruning cache at: [CACHE_DIR]/
+    DEBUG Removing dangling cache entry: [CACHE_DIR]/environments-v1/[ENTRY]
+    Removed [N] files ([SIZE])
+    "###);
+}
+
 /// `cache prune` should remove any stale symlink from the cache.
 #[test]
 fn prune_stale_symlink() -> Result<()> {
diff --git a/crates/uv/tests/common/mod.rs b/crates/uv/tests/common/mod.rs
index 15b33e4a8..e158469ca 100644
--- a/crates/uv/tests/common/mod.rs
+++ b/crates/uv/tests/common/mod.rs
@@ -108,6 +108,10 @@ impl TestContext {
                 format!("{verb} [N] packages"),
             ));
         }
+        self.filters.push((
+            "Removed \\d+ files?".to_string(),
+            "Removed [N] files".to_string(),
+        ));
         self
     }