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