From afb571643f2b8669d12cd09301f9a9f85a8d09d0 Mon Sep 17 00:00:00 2001
From: Charlie Marsh <charlie.r.marsh@gmail.com>
Date: Wed, 24 Jan 2024 10:01:16 -0500
Subject: [PATCH] Avoid unzipping local wheels when fresh (#1076)

Since the archive is a single file in this case, we can rely on the
modification timestamp to check for freshness.
---
 .../src/distribution_database.rs              | 34 +++++++++++++++++--
 crates/puffin-fs/src/lib.rs                   | 17 +++++++---
 2 files changed, 45 insertions(+), 6 deletions(-)

diff --git a/crates/puffin-distribution/src/distribution_database.rs b/crates/puffin-distribution/src/distribution_database.rs
index f8f6aefae..15a389dcc 100644
--- a/crates/puffin-distribution/src/distribution_database.rs
+++ b/crates/puffin-distribution/src/distribution_database.rs
@@ -17,6 +17,7 @@ use platform_tags::Tags;
 use puffin_cache::{Cache, CacheBucket, WheelCache};
 use puffin_client::{CacheControl, CachedClientError, RegistryClient};
 use puffin_extract::unzip_no_seek;
+use puffin_fs::metadata_if_exists;
 use puffin_git::GitSource;
 use puffin_traits::{BuildContext, NoBinary};
 use pypi_types::Metadata21;
@@ -131,7 +132,22 @@ impl<'a, Context: BuildContext + Send + Sync> DistributionDatabase<'a, Context>
                             wheel.filename.stem(),
                         );
 
-                        // TODO(charlie): There's no need to unzip if the wheel is unchanged.
+                        // If the file is already unzipped, and the unzipped directory is fresh,
+                        // return it.
+                        if let (Some(cache_metadata), Some(path_metadata)) = (
+                            metadata_if_exists(cache_entry.path())?,
+                            metadata_if_exists(path)?,
+                        ) {
+                            if cache_metadata.modified()? > path_metadata.modified()? {
+                                return Ok(LocalWheel::Unzipped(UnzippedWheel {
+                                    dist: dist.clone(),
+                                    target: cache_entry.into_path_buf(),
+                                    filename: wheel.filename.clone(),
+                                }));
+                            }
+                        }
+
+                        // Otherwise, unzip the file.
                         return Ok(LocalWheel::Disk(DiskWheel {
                             dist: dist.clone(),
                             path: path.clone(),
@@ -256,7 +272,21 @@ impl<'a, Context: BuildContext + Send + Sync> DistributionDatabase<'a, Context>
                     wheel.filename.stem(),
                 );
 
-                // TODO(charlie): There's no need to unzip if the wheel is unchanged.
+                // If the file is already unzipped, and the unzipped directory is fresh,
+                // return it.
+                if let (Some(cache_metadata), Some(path_metadata)) = (
+                    metadata_if_exists(cache_entry.path())?,
+                    metadata_if_exists(&wheel.path)?,
+                ) {
+                    if cache_metadata.modified()? > path_metadata.modified()? {
+                        return Ok(LocalWheel::Unzipped(UnzippedWheel {
+                            dist: dist.clone(),
+                            target: cache_entry.into_path_buf(),
+                            filename: wheel.filename.clone(),
+                        }));
+                    }
+                }
+
                 Ok(LocalWheel::Disk(DiskWheel {
                     dist: dist.clone(),
                     path: wheel.path.clone(),
diff --git a/crates/puffin-fs/src/lib.rs b/crates/puffin-fs/src/lib.rs
index 9bf205d9d..ff0544130 100644
--- a/crates/puffin-fs/src/lib.rs
+++ b/crates/puffin-fs/src/lib.rs
@@ -68,10 +68,8 @@ pub fn write_atomic_sync(path: impl AsRef<Path>, data: impl AsRef<[u8]>) -> std:
 pub fn force_remove_all(path: impl AsRef<Path>) -> Result<bool, std::io::Error> {
     let path = path.as_ref();
 
-    let metadata = match fs::metadata(path) {
-        Ok(metadata) => metadata,
-        Err(err) if err.kind() == std::io::ErrorKind::NotFound => return Ok(false),
-        Err(err) => return Err(err),
+    let Some(metadata) = metadata_if_exists(path)? else {
+        return Ok(false);
     };
 
     if metadata.is_dir() {
@@ -189,3 +187,14 @@ impl Drop for LockedFile {
         }
     }
 }
+
+/// Given a path, return its metadata if the file exists, or `None` if it does not.
+///
+/// If the file exists but cannot be read, returns an error.
+pub fn metadata_if_exists(path: impl AsRef<Path>) -> std::io::Result<Option<std::fs::Metadata>> {
+    match fs::metadata(path) {
+        Ok(metadata) => Ok(Some(metadata)),
+        Err(err) if err.kind() == std::io::ErrorKind::NotFound => Ok(None),
+        Err(err) => Err(err),
+    }
+}