From ad82b9485683bdbe9e17327ee74cf31b15a16336 Mon Sep 17 00:00:00 2001
From: eth3lbert <eth3lbert+dev@gmail.com>
Date: Tue, 3 Sep 2024 09:20:01 +0800
Subject: [PATCH] Support `file://` URLs for `UV_PYTHON_INSTALL_MIRROR` (#6950)
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

## Summary

Closes #6319.

## Test Plan

I tested with `file:///mirror`, `file://localhost/mirror`, and
`http://mirror` to confirm that it was working as expected.

``` shell-session
/private/tmp/mirror-local                                                                                                                                                                      07:08:18
:)  tree mirror
mirror/
└── 20240814/
   └── cpython-3.12.5+20240814-aarch64-apple-darwin-install_only_stripped.tar.gz
```

<img width="626" alt="image"
src="https://github.com/user-attachments/assets/9c04224d-305c-47ee-a524-4a6abeb79da4">
---
 crates/uv-python/src/downloads.rs | 56 ++++++++++++++++++++++---------
 1 file changed, 40 insertions(+), 16 deletions(-)

diff --git a/crates/uv-python/src/downloads.rs b/crates/uv-python/src/downloads.rs
index af9ef7fe8..6786684a0 100644
--- a/crates/uv-python/src/downloads.rs
+++ b/crates/uv-python/src/downloads.rs
@@ -1,17 +1,17 @@
+use distribution_filename::{ExtensionError, SourceDistExtension};
+use futures::TryStreamExt;
+use owo_colors::OwoColorize;
+use pypi_types::{HashAlgorithm, HashDigest};
 use std::fmt::Display;
 use std::io;
 use std::path::{Path, PathBuf};
 use std::pin::Pin;
 use std::str::FromStr;
 use std::task::{Context, Poll};
-
-use distribution_filename::{ExtensionError, SourceDistExtension};
-use futures::TryStreamExt;
-use owo_colors::OwoColorize;
-use pypi_types::{HashAlgorithm, HashDigest};
 use thiserror::Error;
 use tokio::io::{AsyncRead, ReadBuf};
 use tokio_util::compat::FuturesAsyncReadCompatExt;
+use tokio_util::either::Either;
 use tracing::{debug, instrument};
 use url::Url;
 use uv_client::WrappedReqwestError;
@@ -54,6 +54,8 @@ pub enum Error {
     },
     #[error("Invalid download URL")]
     InvalidUrl(#[from] url::ParseError),
+    #[error("Invalid path in file URL: `{0}`")]
+    InvalidFileUrl(String),
     #[error("Failed to create download directory")]
     DownloadDirError(#[source] io::Error),
     #[error("Failed to copy to: {0}", to.user_display())]
@@ -432,12 +434,8 @@ impl ManagedPythonDownload {
         let filename = url.path_segments().unwrap().last().unwrap();
         let ext = SourceDistExtension::from_path(filename)
             .map_err(|err| Error::MissingExtension(url.to_string(), err))?;
-        let response = client.client().get(url.clone()).send().await?;
+        let (reader, size) = read_url(&url, client).await?;
 
-        // Ensure the request was successful.
-        response.error_for_status_ref()?;
-
-        let size = response.content_length();
         let progress = reporter
             .as_ref()
             .map(|reporter| (reporter, reporter.on_download_start(&self.key, size)));
@@ -450,17 +448,12 @@ impl ManagedPythonDownload {
             temp_dir.path().simplified().display()
         );
 
-        let stream = response
-            .bytes_stream()
-            .map_err(|err| io::Error::new(io::ErrorKind::Other, err))
-            .into_async_read();
-
         let mut hashers = self
             .sha256
             .into_iter()
             .map(|_| Hasher::from(HashAlgorithm::Sha256))
             .collect::<Vec<_>>();
-        let mut hasher = uv_extract::hash::HashReader::new(stream.compat(), &mut hashers);
+        let mut hasher = uv_extract::hash::HashReader::new(reader, &mut hashers);
 
         debug!("Extracting {filename}");
 
@@ -640,3 +633,34 @@ where
             })
     }
 }
+
+/// Convert a [`Url`] into an [`AsyncRead`] stream.
+async fn read_url(
+    url: &Url,
+    client: &uv_client::BaseClient,
+) -> Result<(impl AsyncRead + Unpin, Option<u64>), Error> {
+    if url.scheme() == "file" {
+        // Loads downloaded distribution from the given `file://` URL.
+        let path = url
+            .to_file_path()
+            .map_err(|()| Error::InvalidFileUrl(url.to_string()))?;
+
+        let size = fs_err::tokio::metadata(&path).await?.len();
+        let reader = fs_err::tokio::File::open(&path).await?;
+
+        Ok((Either::Left(reader), Some(size)))
+    } else {
+        let response = client.client().get(url.clone()).send().await?;
+
+        // Ensure the request was successful.
+        response.error_for_status_ref()?;
+
+        let size = response.content_length();
+        let stream = response
+            .bytes_stream()
+            .map_err(|err| io::Error::new(io::ErrorKind::Other, err))
+            .into_async_read();
+
+        Ok((Either::Right(stream.compat()), size))
+    }
+}