Skip to content

Commit

Permalink
Support file:// URLs for UV_PYTHON_INSTALL_MIRROR (#6950)
Browse files Browse the repository at this point in the history
## 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">
  • Loading branch information
eth3lbert committed Sep 3, 2024
1 parent b229230 commit ad82b94
Showing 1 changed file with 40 additions and 16 deletions.
56 changes: 40 additions & 16 deletions crates/uv-python/src/downloads.rs
Original file line number Diff line number Diff line change
@@ -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;
Expand Down Expand Up @@ -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())]
Expand Down Expand Up @@ -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)));
Expand All @@ -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}");

Expand Down Expand Up @@ -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))
}
}

0 comments on commit ad82b94

Please sign in to comment.