diff --git a/crates/puffin-cli/src/commands/pip_compile.rs b/crates/puffin-cli/src/commands/pip_compile.rs index 43fbab8d2..1047a3502 100644 --- a/crates/puffin-cli/src/commands/pip_compile.rs +++ b/crates/puffin-cli/src/commands/pip_compile.rs @@ -26,7 +26,7 @@ use puffin_normalize::ExtraName; use puffin_resolver::{ DisplayResolutionGraph, Manifest, PreReleaseMode, ResolutionMode, ResolutionOptions, Resolver, }; -use puffin_traits::SetupPyStrategy; +use puffin_traits::{InFlight, SetupPyStrategy}; use requirements_txt::EditableRequirement; use crate::commands::reporters::{DownloadReporter, ResolverReporter}; @@ -154,6 +154,9 @@ pub(crate) async fn pip_compile( FlatIndex::from_entries(entries, &tags) }; + // Track in-flight downloads, builds, etc., across resolutions. + let in_flight = InFlight::default(); + let options = ResolutionOptions::new(resolution_mode, prerelease_mode, exclude_newer); let build_dispatch = BuildDispatch::new( &client, @@ -161,6 +164,7 @@ pub(crate) async fn pip_compile( &interpreter, &index_locations, &flat_index, + &in_flight, interpreter.sys_executable().to_path_buf(), setup_py, no_build, @@ -227,13 +231,6 @@ pub(crate) async fn pip_compile( editable_metadata, ); - // Resolve the flat indexes from `--find-links`. - let flat_index = { - let client = FlatIndexClient::new(&client, &cache); - let entries = client.fetch(index_locations.flat_indexes()).await?; - FlatIndex::from_entries(entries, &tags) - }; - // Resolve the dependencies. let resolver = Resolver::new( manifest, diff --git a/crates/puffin-cli/src/commands/pip_install.rs b/crates/puffin-cli/src/commands/pip_install.rs index 8d89f171d..27a995182 100644 --- a/crates/puffin-cli/src/commands/pip_install.rs +++ b/crates/puffin-cli/src/commands/pip_install.rs @@ -28,7 +28,7 @@ use puffin_normalize::PackageName; use puffin_resolver::{ Manifest, PreReleaseMode, ResolutionGraph, ResolutionMode, ResolutionOptions, Resolver, }; -use puffin_traits::{OnceMap, SetupPyStrategy}; +use puffin_traits::{InFlight, SetupPyStrategy}; use requirements_txt::EditableRequirement; use crate::commands::reporters::{DownloadReporter, InstallReporter, ResolverReporter}; @@ -144,6 +144,9 @@ pub(crate) async fn pip_install( FlatIndex::from_entries(entries, tags) }; + // Track in-flight downloads, builds, etc., across resolutions. + let in_flight = InFlight::default(); + let options = ResolutionOptions::new(resolution_mode, prerelease_mode, exclude_newer); let build_dispatch = BuildDispatch::new( @@ -152,6 +155,7 @@ pub(crate) async fn pip_install( &interpreter, &index_locations, &flat_index, + &in_flight, venv.python_executable(), setup_py, no_build, @@ -221,6 +225,7 @@ pub(crate) async fn pip_install( &index_locations, tags, &client, + &in_flight, &build_dispatch, &cache, &venv, @@ -420,6 +425,7 @@ async fn install( index_urls: &IndexLocations, tags: &Tags, client: &RegistryClient, + in_flight: &InFlight, build_dispatch: &BuildDispatch<'_>, cache: &Cache, venv: &Virtualenv, @@ -488,8 +494,9 @@ async fn install( let downloader = Downloader::new(cache, tags, client, build_dispatch) .with_reporter(DownloadReporter::from(printer).with_length(remote.len() as u64)); + // STOPSHIP(charlie): This needs to be shared! let wheels = downloader - .download(remote, &OnceMap::default()) + .download(remote, in_flight) .await .context("Failed to download distributions")?; diff --git a/crates/puffin-cli/src/commands/pip_sync.rs b/crates/puffin-cli/src/commands/pip_sync.rs index 983ebf56b..a2d5352ab 100644 --- a/crates/puffin-cli/src/commands/pip_sync.rs +++ b/crates/puffin-cli/src/commands/pip_sync.rs @@ -14,7 +14,7 @@ use puffin_client::{FlatIndex, FlatIndexClient, RegistryClient, RegistryClientBu use puffin_dispatch::BuildDispatch; use puffin_installer::{Downloader, InstallPlan, Reinstall, ResolvedEditable, SitePackages}; use puffin_interpreter::Virtualenv; -use puffin_traits::{OnceMap, SetupPyStrategy}; +use puffin_traits::{InFlight, SetupPyStrategy}; use pypi_types::Yanked; use requirements_txt::EditableRequirement; @@ -70,6 +70,9 @@ pub(crate) async fn pip_sync( FlatIndex::from_entries(entries, tags) }; + // Track in-flight downloads, builds, etc., across resolutions. + let in_flight = InFlight::default(); + // Prep the build context. let build_dispatch = BuildDispatch::new( &client, @@ -77,6 +80,7 @@ pub(crate) async fn pip_sync( venv.interpreter(), &index_locations, &flat_index, + &in_flight, venv.python_executable(), setup_py, no_build, @@ -211,7 +215,7 @@ pub(crate) async fn pip_sync( .with_reporter(DownloadReporter::from(printer).with_length(remote.len() as u64)); let wheels = downloader - .download(remote, &OnceMap::default()) + .download(remote, &in_flight) .await .context("Failed to download distributions")?; diff --git a/crates/puffin-cli/src/commands/venv.rs b/crates/puffin-cli/src/commands/venv.rs index cb668183f..14dbf332c 100644 --- a/crates/puffin-cli/src/commands/venv.rs +++ b/crates/puffin-cli/src/commands/venv.rs @@ -15,7 +15,7 @@ use puffin_cache::Cache; use puffin_client::{FlatIndex, FlatIndexClient, RegistryClientBuilder}; use puffin_dispatch::BuildDispatch; use puffin_interpreter::Interpreter; -use puffin_traits::{BuildContext, SetupPyStrategy}; +use puffin_traits::{BuildContext, InFlight, SetupPyStrategy}; use crate::commands::ExitStatus; use crate::printer::Printer; @@ -139,6 +139,9 @@ async fn venv_impl( FlatIndex::from_entries(entries, tags) }; + // Track in-flight downloads, builds, etc., across resolutions. + let in_flight = InFlight::default(); + // Prep the build context. let build_dispatch = BuildDispatch::new( &client, @@ -146,6 +149,7 @@ async fn venv_impl( interpreter, index_locations, &flat_index, + &in_flight, venv.python_executable(), SetupPyStrategy::default(), true, diff --git a/crates/puffin-dev/src/build.rs b/crates/puffin-dev/src/build.rs index fc6748245..f2d0789e1 100644 --- a/crates/puffin-dev/src/build.rs +++ b/crates/puffin-dev/src/build.rs @@ -12,7 +12,7 @@ use puffin_cache::{Cache, CacheArgs}; use puffin_client::{FlatIndex, RegistryClientBuilder}; use puffin_dispatch::BuildDispatch; use puffin_interpreter::Virtualenv; -use puffin_traits::{BuildContext, BuildKind, SetupPyStrategy}; +use puffin_traits::{BuildContext, BuildKind, InFlight, SetupPyStrategy}; #[derive(Parser)] pub(crate) struct BuildArgs { @@ -57,6 +57,7 @@ pub(crate) async fn build(args: BuildArgs) -> Result { let index_urls = IndexLocations::default(); let flat_index = FlatIndex::default(); let setup_py = SetupPyStrategy::default(); + let in_flight = InFlight::default(); let build_dispatch = BuildDispatch::new( &client, @@ -64,6 +65,7 @@ pub(crate) async fn build(args: BuildArgs) -> Result { venv.interpreter(), &index_urls, &flat_index, + &in_flight, venv.python_executable(), setup_py, false, diff --git a/crates/puffin-dev/src/install_many.rs b/crates/puffin-dev/src/install_many.rs index e235ffc8e..0d3aaa3f4 100644 --- a/crates/puffin-dev/src/install_many.rs +++ b/crates/puffin-dev/src/install_many.rs @@ -25,7 +25,7 @@ use puffin_installer::Downloader; use puffin_interpreter::Virtualenv; use puffin_normalize::PackageName; use puffin_resolver::DistFinder; -use puffin_traits::{BuildContext, OnceMap, SetupPyStrategy}; +use puffin_traits::{BuildContext, InFlight, SetupPyStrategy}; #[derive(Parser)] pub(crate) struct InstallManyArgs { @@ -62,6 +62,7 @@ pub(crate) async fn install_many(args: InstallManyArgs) -> Result<()> { let index_locations = IndexLocations::default(); let flat_index = FlatIndex::default(); let setup_py = SetupPyStrategy::default(); + let in_flight = InFlight::default(); let tags = venv.interpreter().tags()?; let build_dispatch = BuildDispatch::new( @@ -70,6 +71,7 @@ pub(crate) async fn install_many(args: InstallManyArgs) -> Result<()> { venv.interpreter(), &index_locations, &flat_index, + &in_flight, venv.python_executable(), setup_py, args.no_build, @@ -141,7 +143,7 @@ async fn install_chunk( let mut registry_index = RegistryWheelIndex::new(build_dispatch.cache(), tags, index_locations); let (cached, uncached): (Vec<_>, Vec<_>) = dists.into_iter().partition_map(|dist| { - // We always want the wheel for the latest version not whatever matching is in cache + // We always want the wheel for the latest version not whatever matching is in cache. let VersionOrUrl::Version(version) = dist.version_or_url() else { unreachable!(); }; @@ -155,7 +157,7 @@ async fn install_chunk( info!("Cached: {}, Uncached {}", cached.len(), uncached.len()); let downloader = Downloader::new(build_dispatch.cache(), tags, client, build_dispatch); - let in_flight = OnceMap::default(); + let in_flight = InFlight::default(); let fetches: Vec<_> = futures::stream::iter(uncached) .map(|dist| downloader.get_wheel(dist, &in_flight)) .buffer_unordered(50) diff --git a/crates/puffin-dev/src/resolve_cli.rs b/crates/puffin-dev/src/resolve_cli.rs index 3abe97c66..3c566395c 100644 --- a/crates/puffin-dev/src/resolve_cli.rs +++ b/crates/puffin-dev/src/resolve_cli.rs @@ -17,7 +17,7 @@ use puffin_client::{FlatIndex, FlatIndexClient, RegistryClientBuilder}; use puffin_dispatch::BuildDispatch; use puffin_interpreter::Virtualenv; use puffin_resolver::{Manifest, ResolutionOptions, Resolver}; -use puffin_traits::SetupPyStrategy; +use puffin_traits::{InFlight, SetupPyStrategy}; #[derive(ValueEnum, Default, Clone)] pub(crate) enum ResolveCliFormat { @@ -65,6 +65,7 @@ pub(crate) async fn resolve_cli(args: ResolveCliArgs) -> Result<()> { let entries = client.fetch(index_locations.flat_indexes()).await?; FlatIndex::from_entries(entries, venv.interpreter().tags()?) }; + let in_flight = InFlight::default(); let build_dispatch = BuildDispatch::new( &client, @@ -72,6 +73,7 @@ pub(crate) async fn resolve_cli(args: ResolveCliArgs) -> Result<()> { venv.interpreter(), &index_locations, &flat_index, + &in_flight, venv.python_executable(), SetupPyStrategy::default(), args.no_build, diff --git a/crates/puffin-dev/src/resolve_many.rs b/crates/puffin-dev/src/resolve_many.rs index 9a9ab60a6..ad6b0f551 100644 --- a/crates/puffin-dev/src/resolve_many.rs +++ b/crates/puffin-dev/src/resolve_many.rs @@ -20,7 +20,7 @@ use puffin_client::{FlatIndex, RegistryClient, RegistryClientBuilder}; use puffin_dispatch::BuildDispatch; use puffin_interpreter::Virtualenv; use puffin_normalize::PackageName; -use puffin_traits::{BuildContext, SetupPyStrategy}; +use puffin_traits::{BuildContext, InFlight, SetupPyStrategy}; #[derive(Parser)] pub(crate) struct ResolveManyArgs { @@ -76,6 +76,7 @@ pub(crate) async fn resolve_many(args: ResolveManyArgs) -> Result<()> { let index_locations = IndexLocations::default(); let flat_index = FlatIndex::default(); let setup_py = SetupPyStrategy::default(); + let in_flight = InFlight::default(); let build_dispatch = BuildDispatch::new( &client, @@ -83,6 +84,7 @@ pub(crate) async fn resolve_many(args: ResolveManyArgs) -> Result<()> { venv.interpreter(), &index_locations, &flat_index, + &in_flight, venv.python_executable(), setup_py, args.no_build, diff --git a/crates/puffin-dispatch/src/lib.rs b/crates/puffin-dispatch/src/lib.rs index fc67d24bf..554575f8c 100644 --- a/crates/puffin-dispatch/src/lib.rs +++ b/crates/puffin-dispatch/src/lib.rs @@ -9,7 +9,7 @@ use anyhow::{bail, Context, Result}; use itertools::Itertools; use tracing::{debug, instrument}; -use distribution_types::{CachedDist, DistributionId, IndexLocations, Name, Resolution}; +use distribution_types::{IndexLocations, Name, Resolution}; use pep508_rs::Requirement; use puffin_build::{SourceBuild, SourceBuildContext}; use puffin_cache::Cache; @@ -17,7 +17,7 @@ use puffin_client::{FlatIndex, RegistryClient}; use puffin_installer::{Downloader, InstallPlan, Installer, Reinstall, SitePackages}; use puffin_interpreter::{Interpreter, Virtualenv}; use puffin_resolver::{Manifest, ResolutionOptions, Resolver}; -use puffin_traits::{BuildContext, BuildKind, OnceMap, SetupPyStrategy}; +use puffin_traits::{BuildContext, BuildKind, InFlight, SetupPyStrategy}; /// The main implementation of [`BuildContext`], used by the CLI, see [`BuildContext`] /// documentation. @@ -27,12 +27,12 @@ pub struct BuildDispatch<'a> { interpreter: &'a Interpreter, index_locations: &'a IndexLocations, flat_index: &'a FlatIndex, + in_flight: &'a InFlight, base_python: PathBuf, setup_py: SetupPyStrategy, no_build: bool, source_build_context: SourceBuildContext, options: ResolutionOptions, - in_flight: OnceMap>, } impl<'a> BuildDispatch<'a> { @@ -43,6 +43,7 @@ impl<'a> BuildDispatch<'a> { interpreter: &'a Interpreter, index_locations: &'a IndexLocations, flat_index: &'a FlatIndex, + in_flight: &'a InFlight, base_python: PathBuf, setup_py: SetupPyStrategy, no_build: bool, @@ -53,12 +54,12 @@ impl<'a> BuildDispatch<'a> { interpreter, index_locations, flat_index, + in_flight, base_python, setup_py, no_build, source_build_context: SourceBuildContext::default(), options: ResolutionOptions::default(), - in_flight: OnceMap::default(), } } @@ -184,7 +185,7 @@ impl<'a> BuildContext for BuildDispatch<'a> { ); downloader - .download(remote, &self.in_flight) + .download(remote, self.in_flight) .await .context("Failed to download and build distributions")? }; diff --git a/crates/puffin-installer/src/downloader.rs b/crates/puffin-installer/src/downloader.rs index f46f0a183..ca716adc8 100644 --- a/crates/puffin-installer/src/downloader.rs +++ b/crates/puffin-installer/src/downloader.rs @@ -7,14 +7,12 @@ use tokio::task::JoinError; use tracing::{instrument, warn}; use url::Url; -use distribution_types::{ - CachedDist, Dist, DistributionId, Identifier, LocalEditable, RemoteSource, SourceDist, -}; +use distribution_types::{CachedDist, Dist, Identifier, LocalEditable, RemoteSource, SourceDist}; use platform_tags::Tags; use puffin_cache::Cache; use puffin_client::RegistryClient; use puffin_distribution::{DistributionDatabase, DistributionDatabaseError, LocalWheel, Unzip}; -use puffin_traits::{BuildContext, OnceMap}; +use puffin_traits::{BuildContext, InFlight}; use crate::editable::BuiltEditable; @@ -66,7 +64,7 @@ impl<'a, Context: BuildContext + Send + Sync> Downloader<'a, Context> { pub fn download_stream<'stream>( &'stream self, distributions: Vec, - in_flight: &'stream OnceMap>, + in_flight: &'stream InFlight, ) -> impl Stream> + 'stream { futures::stream::iter(distributions) .map(|dist| async { @@ -86,7 +84,7 @@ impl<'a, Context: BuildContext + Send + Sync> Downloader<'a, Context> { pub async fn download( &self, mut distributions: Vec, - in_flight: &OnceMap>, + in_flight: &InFlight, ) -> Result, Error> { // Sort the distributions by size. distributions @@ -154,13 +152,9 @@ impl<'a, Context: BuildContext + Send + Sync> Downloader<'a, Context> { /// Download, build, and unzip a single wheel. #[instrument(skip_all, fields(name = % dist, size = ? dist.size(), url = dist.file().map(|file| file.url.to_string()).unwrap_or_default()))] - pub async fn get_wheel( - &self, - dist: Dist, - in_flight: &OnceMap>, - ) -> Result { + pub async fn get_wheel(&self, dist: Dist, in_flight: &InFlight) -> Result { let id = dist.distribution_id(); - let wheel = if in_flight.register(&id) { + let wheel = if in_flight.downloads.register(&id) { let download: LocalWheel = self .database .get_or_build_wheel(dist.clone()) @@ -169,16 +163,17 @@ impl<'a, Context: BuildContext + Send + Sync> Downloader<'a, Context> { let result = Self::unzip_wheel(download).await; match result { Ok(cached) => { - in_flight.done(id, Ok(cached.clone())); + in_flight.downloads.done(id, Ok(cached.clone())); cached } Err(err) => { - in_flight.done(id, Err(err.to_string())); + in_flight.downloads.done(id, Err(err.to_string())); return Err(err); } } } else { in_flight + .downloads .wait(&id) .await .value() diff --git a/crates/puffin-traits/src/lib.rs b/crates/puffin-traits/src/lib.rs index e8e5e9f2a..20637697c 100644 --- a/crates/puffin-traits/src/lib.rs +++ b/crates/puffin-traits/src/lib.rs @@ -6,7 +6,7 @@ use std::path::{Path, PathBuf}; use anyhow::Result; -use distribution_types::Resolution; +use distribution_types::{CachedDist, DistributionId, Resolution}; pub use once_map::OnceMap; use pep508_rs::Requirement; use puffin_cache::Cache; @@ -124,6 +124,12 @@ pub trait SourceBuildTrait { -> impl Future> + Send + 'a; } +#[derive(Default)] +pub struct InFlight { + /// The in-flight distribution downloads. + pub downloads: OnceMap>, +} + /// The strategy to use when building source distributions that lack a `pyproject.toml`. #[derive(Copy, Clone, Debug, Default, PartialEq, Eq)] pub enum SetupPyStrategy {