From c7fc25a0e75536f1bfae172e909a726ee78594ae Mon Sep 17 00:00:00 2001
From: John Sirois <john.sirois@gmail.com>
Date: Mon, 7 Oct 2024 10:45:35 -0700
Subject: [PATCH] Fix prune.

---
 pex/cache/data.py                 |  6 +++---
 pex/cli/commands/cache/command.py | 22 +++++++++++-----------
 2 files changed, 14 insertions(+), 14 deletions(-)

diff --git a/pex/cache/data.py b/pex/cache/data.py
index a96a375fb..5f25c4db2 100644
--- a/pex/cache/data.py
+++ b/pex/cache/data.py
@@ -348,7 +348,7 @@ def prune(
                     keys,
                 )
             ) as cursor:
-                for bootstrap_hash in cursor:
+                for [bootstrap_hash] in cursor:
                     bootstraps_by_hash.pop(bootstrap_hash)
 
         user_code_by_hash = OrderedDict(
@@ -365,7 +365,7 @@ def prune(
                     keys,
                 )
             ) as cursor:
-                for code_hash in cursor:
+                for [code_hash] in cursor:
                     user_code_by_hash.pop(code_hash)
 
         wheels_by_hash = OrderedDict(
@@ -388,7 +388,7 @@ def prune(
                     keys,
                 )
             ) as cursor:
-                for install_hash in cursor:
+                for [install_hash] in cursor:
                     wheels_by_hash.pop(install_hash)
 
         yield itertools.chain(
diff --git a/pex/cli/commands/cache/command.py b/pex/cli/commands/cache/command.py
index bb1e76c2f..4cf7bd536 100644
--- a/pex/cli/commands/cache/command.py
+++ b/pex/cli/commands/cache/command.py
@@ -10,7 +10,7 @@
 
 from pex.cache import access as cache_access
 from pex.cache import data as cache_data
-from pex.cache.dirs import AtomicCacheDir, CacheDir
+from pex.cache.dirs import CacheDir
 from pex.cli.command import BuildTimeCommand
 from pex.cli.commands.cache.bytes import ByteAmount, ByteUnits
 from pex.cli.commands.cache.du import DiskUsage
@@ -441,11 +441,11 @@ def _purge(self):
 
         return Ok()
 
-    def _prune_atomic_cache_dir(self, atomic_cache_dir):
-        # type: (AtomicCacheDir) -> DiskUsage
-        du = DiskUsage.collect(atomic_cache_dir.path)
+    def _prune_cache_dir(self, cache_dir):
+        # type: (str) -> DiskUsage
+        du = DiskUsage.collect(cache_dir)
         if not self.options.dry_run:
-            safe_rmtree(atomic_cache_dir.path)
+            safe_rmtree(cache_dir)
         return du
 
     def _prune(self):
@@ -495,8 +495,8 @@ def _prune(self):
             self._render_usage(
                 list(
                     iter_map_parallel(
-                        pex_dirs,
-                        self._prune_atomic_cache_dir,
+                        (pex_dir.path for pex_dir in pex_dirs),
+                        self._prune_cache_dir,
                         noun="cached PEX",
                         verb="prune",
                         verb_past="pruned",
@@ -518,8 +518,8 @@ def _prune(self):
                 self._render_usage(
                     list(
                         iter_map_parallel(
-                            deps,
-                            self._prune_atomic_cache_dir,
+                            (dep.path for dep in deps),
+                            self._prune_cache_dir,
                             noun="cached PEX dependency",
                             verb="prune",
                             verb_past="pruned",
@@ -532,8 +532,8 @@ def _prune(self):
             with cache_data.prune(deps) as prunable_deps_iter:
                 disk_usages = list(
                     iter_map_parallel(
-                        prunable_deps_iter,
-                        self._prune_atomic_cache_dir,
+                        (dep.path for dep in prunable_deps_iter),
+                        self._prune_cache_dir,
                         noun="cached PEX dependency",
                         verb="prune",
                         verb_past="pruned",