Do wheel downloads concurrently (#28)

This commit is contained in:
Charlie Marsh 2023-10-06 16:51:31 -04:00 committed by GitHub
parent dd26cfa0cc
commit 36d0124e60
No known key found for this signature in database
GPG key ID: 4AEE18F83AFDEB23

View file

@ -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)
}