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 std::str::FromStr;
|
||||||
|
|
||||||
use anyhow::Result;
|
use anyhow::Result;
|
||||||
use install_wheel_rs::{install_wheel, InstallLocation};
|
use install_wheel_rs::{install_wheel, InstallLocation};
|
||||||
|
use tokio::task::JoinSet;
|
||||||
use tokio_util::compat::FuturesAsyncReadCompatExt;
|
use tokio_util::compat::FuturesAsyncReadCompatExt;
|
||||||
use url::Url;
|
use url::Url;
|
||||||
|
|
||||||
|
@ -17,20 +19,20 @@ pub async fn install(
|
||||||
// Create a temporary directory, in which we'll store the wheels.
|
// Create a temporary directory, in which we'll store the wheels.
|
||||||
let tmp_dir = tempfile::tempdir()?;
|
let tmp_dir = tempfile::tempdir()?;
|
||||||
|
|
||||||
// Download each wheel.
|
// Download the wheels in parallel.
|
||||||
// TODO(charlie): Store these in a content-addressed cache.
|
let mut downloads = JoinSet::new();
|
||||||
// TODO(charlie): Use channels to efficiently stream-and-install.
|
|
||||||
for wheel in wheels {
|
for wheel in wheels {
|
||||||
let url = Url::parse(&wheel.url)?;
|
downloads.spawn(do_download(
|
||||||
let reader = client.stream_external(&url).await?;
|
wheel.clone(),
|
||||||
|
client.clone(),
|
||||||
// TODO(charlie): Stream the unzip.
|
tmp_dir.path().join(&wheel.hashes.sha256),
|
||||||
let mut writer = tokio::fs::File::create(tmp_dir.path().join(&wheel.hashes.sha256)).await?;
|
));
|
||||||
tokio::io::copy(&mut reader.compat(), &mut writer).await?;
|
}
|
||||||
|
while let Some(result) = downloads.join_next().await.transpose()? {
|
||||||
|
result?;
|
||||||
}
|
}
|
||||||
|
|
||||||
// Install each wheel.
|
// Install each wheel.
|
||||||
// TODO(charlie): Use channels to efficiently stream-and-install.
|
|
||||||
let location = InstallLocation::Venv {
|
let location = InstallLocation::Venv {
|
||||||
venv_base: python.venv().to_path_buf(),
|
venv_base: python.venv().to_path_buf(),
|
||||||
python_version: python.simple_version(),
|
python_version: python.simple_version(),
|
||||||
|
@ -54,3 +56,16 @@ pub async fn install(
|
||||||
|
|
||||||
Ok(())
|
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