mirror of
https://github.com/astral-sh/uv.git
synced 2025-07-07 13:25:00 +00:00
Do wheel downloads concurrently (#28)
This commit is contained in:
parent
dd26cfa0cc
commit
36d0124e60
1 changed files with 25 additions and 10 deletions
|
@ -1,7 +1,9 @@
|
|||
use std::path::Path;
|
||||
use std::str::FromStr;
|
||||
|
||||
use anyhow::Result;
|
||||
use install_wheel_rs::{install_wheel, InstallLocation};
|
||||
use tokio::task::JoinSet;
|
||||
use tokio_util::compat::FuturesAsyncReadCompatExt;
|
||||
use url::Url;
|
||||
|
||||
|
@ -17,20 +19,20 @@ pub async fn install(
|
|||
// Create a temporary directory, in which we'll store the wheels.
|
||||
let tmp_dir = tempfile::tempdir()?;
|
||||
|
||||
// Download each wheel.
|
||||
// TODO(charlie): Store these in a content-addressed cache.
|
||||
// TODO(charlie): Use channels to efficiently stream-and-install.
|
||||
// Download the wheels in parallel.
|
||||
let mut downloads = JoinSet::new();
|
||||
for wheel in wheels {
|
||||
let url = Url::parse(&wheel.url)?;
|
||||
let reader = client.stream_external(&url).await?;
|
||||
|
||||
// TODO(charlie): Stream the unzip.
|
||||
let mut writer = tokio::fs::File::create(tmp_dir.path().join(&wheel.hashes.sha256)).await?;
|
||||
tokio::io::copy(&mut reader.compat(), &mut writer).await?;
|
||||
downloads.spawn(do_download(
|
||||
wheel.clone(),
|
||||
client.clone(),
|
||||
tmp_dir.path().join(&wheel.hashes.sha256),
|
||||
));
|
||||
}
|
||||
while let Some(result) = downloads.join_next().await.transpose()? {
|
||||
result?;
|
||||
}
|
||||
|
||||
// Install each wheel.
|
||||
// TODO(charlie): Use channels to efficiently stream-and-install.
|
||||
let location = InstallLocation::Venv {
|
||||
venv_base: python.venv().to_path_buf(),
|
||||
python_version: python.simple_version(),
|
||||
|
@ -54,3 +56,16 @@ pub async fn install(
|
|||
|
||||
Ok(())
|
||||
}
|
||||
|
||||
/// Download a wheel to a given path.
|
||||
async fn do_download(wheel: File, client: PypiClient, path: impl AsRef<Path>) -> Result<File> {
|
||||
// TODO(charlie): Store these in a content-addressed cache.
|
||||
let url = Url::parse(&wheel.url)?;
|
||||
let reader = client.stream_external(&url).await?;
|
||||
|
||||
// TODO(charlie): Stream the unzip.
|
||||
let mut writer = tokio::fs::File::create(path).await?;
|
||||
tokio::io::copy(&mut reader.compat(), &mut writer).await?;
|
||||
|
||||
Ok(wheel)
|
||||
}
|
||||
|
|
Loading…
Add table
Add a link
Reference in a new issue