diff --git a/crates/uv-python/src/downloads.rs b/crates/uv-python/src/downloads.rs index af9ef7fe830c..6786684a0ed6 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?; - - // Ensure the request was successful. - response.error_for_status_ref()?; + let (reader, size) = read_url(&url, client).await?; - 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::>(); - 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), 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)) + } +}