From 1637f1c2169e3ca872c24ac45caadd7a0aa4d9d1 Mon Sep 17 00:00:00 2001 From: Charlie Marsh Date: Sun, 5 Nov 2023 21:16:04 -0800 Subject: [PATCH] Add source distribution support to the `DistributionFinder` (#322) ## Summary This just enables the `DistributionFinder` (previously known as the `WheelFinder`) to select source distributions when there are no matching wheels for a given platform. As a reminder, the `DistributionFinder` is a simple resolver that doesn't look at any dependencies: it just takes a set of pinned packages, and finds a distribution to install to satisfy each requirement. --- crates/puffin-cli/src/commands/pip_sync.rs | 6 +- crates/puffin-cli/src/commands/reporters.rs | 8 +- crates/puffin-dispatch/src/lib.rs | 4 +- .../src/{wheel_finder.rs => finder.rs} | 98 +++++++++---------- crates/puffin-resolver/src/lib.rs | 4 +- 5 files changed, 56 insertions(+), 64 deletions(-) rename crates/puffin-resolver/src/{wheel_finder.rs => finder.rs} (62%) diff --git a/crates/puffin-cli/src/commands/pip_sync.rs b/crates/puffin-cli/src/commands/pip_sync.rs index 0b170b288..7f89fd8d5 100644 --- a/crates/puffin-cli/src/commands/pip_sync.rs +++ b/crates/puffin-cli/src/commands/pip_sync.rs @@ -16,7 +16,7 @@ use puffin_installer::PartitionedRequirements; use puffin_interpreter::Virtualenv; use crate::commands::reporters::{ - DownloadReporter, InstallReporter, UnzipReporter, WheelFinderReporter, + DownloadReporter, FinderReporter, InstallReporter, UnzipReporter, }; use crate::commands::{elapsed, ExitStatus}; use crate::index_urls::IndexUrls; @@ -118,8 +118,8 @@ pub(crate) async fn sync_requirements( } else { let start = std::time::Instant::now(); - let wheel_finder = puffin_resolver::WheelFinder::new(&tags, &client) - .with_reporter(WheelFinderReporter::from(printer).with_length(remote.len() as u64)); + let wheel_finder = puffin_resolver::DistributionFinder::new(&tags, &client) + .with_reporter(FinderReporter::from(printer).with_length(remote.len() as u64)); let resolution = wheel_finder.resolve(&remote).await?; let s = if resolution.len() == 1 { "" } else { "s" }; diff --git a/crates/puffin-cli/src/commands/reporters.rs b/crates/puffin-cli/src/commands/reporters.rs index 3acc602ab..ef7e0381f 100644 --- a/crates/puffin-cli/src/commands/reporters.rs +++ b/crates/puffin-cli/src/commands/reporters.rs @@ -9,11 +9,11 @@ use puffin_normalize::PackageName; use crate::printer::Printer; #[derive(Debug)] -pub(crate) struct WheelFinderReporter { +pub(crate) struct FinderReporter { progress: ProgressBar, } -impl From for WheelFinderReporter { +impl From for FinderReporter { fn from(printer: Printer) -> Self { let progress = ProgressBar::with_draw_target(None, printer.target()); progress.set_style( @@ -24,7 +24,7 @@ impl From for WheelFinderReporter { } } -impl WheelFinderReporter { +impl FinderReporter { #[must_use] pub(crate) fn with_length(self, length: u64) -> Self { self.progress.set_length(length); @@ -32,7 +32,7 @@ impl WheelFinderReporter { } } -impl puffin_resolver::WheelFinderReporter for WheelFinderReporter { +impl puffin_resolver::FinderReporter for FinderReporter { fn on_progress(&self, wheel: &RemoteDistribution) { self.progress.set_message(format!("{wheel}")); self.progress.inc(1); diff --git a/crates/puffin-dispatch/src/lib.rs b/crates/puffin-dispatch/src/lib.rs index 86cd10e32..5ab815f0e 100644 --- a/crates/puffin-dispatch/src/lib.rs +++ b/crates/puffin-dispatch/src/lib.rs @@ -17,7 +17,7 @@ use puffin_build::SourceDistributionBuilder; use puffin_client::RegistryClient; use puffin_installer::{Downloader, Installer, PartitionedRequirements, Unzipper}; use puffin_interpreter::{InterpreterInfo, Virtualenv}; -use puffin_resolver::{Manifest, PreReleaseMode, ResolutionMode, Resolver, WheelFinder}; +use puffin_resolver::{DistributionFinder, Manifest, PreReleaseMode, ResolutionMode, Resolver}; use puffin_traits::BuildContext; /// The main implementation of [`BuildContext`], used by the CLI, see [`BuildContext`] @@ -122,7 +122,7 @@ impl BuildContext for BuildDispatch { if remote.len() == 1 { "" } else { "s" }, remote.iter().map(ToString::to_string).join(", ") ); - let resolution = WheelFinder::new(&tags, &self.client) + let resolution = DistributionFinder::new(&tags, &self.client) .resolve(&remote) .await .context("Failed to resolve build dependencies")?; diff --git a/crates/puffin-resolver/src/wheel_finder.rs b/crates/puffin-resolver/src/finder.rs similarity index 62% rename from crates/puffin-resolver/src/wheel_finder.rs rename to crates/puffin-resolver/src/finder.rs index 893dc87d2..4375f2392 100644 --- a/crates/puffin-resolver/src/wheel_finder.rs +++ b/crates/puffin-resolver/src/finder.rs @@ -1,4 +1,4 @@ -//! Given a set of selected packages, find a compatible set of wheels to install. +//! Given a set of selected packages, find a compatible set of distributions to install. //! //! This is similar to running `pip install` with the `--no-deps` flag. @@ -6,30 +6,28 @@ use std::hash::BuildHasherDefault; use std::str::FromStr; use anyhow::Result; -use futures::future::Either; use futures::{StreamExt, TryFutureExt}; use fxhash::FxHashMap; -use tracing::debug; -use distribution_filename::WheelFilename; +use distribution_filename::{SourceDistributionFilename, WheelFilename}; use pep508_rs::{Requirement, VersionOrUrl}; use platform_tags::Tags; use puffin_client::RegistryClient; use puffin_distribution::RemoteDistribution; use puffin_normalize::PackageName; -use puffin_package::pypi_types::{File, Metadata21, SimpleJson}; +use puffin_package::pypi_types::{File, SimpleJson}; use crate::error::ResolveError; use crate::resolution::Resolution; -pub struct WheelFinder<'a> { +pub struct DistributionFinder<'a> { tags: &'a Tags, client: &'a RegistryClient, reporter: Option>, } -impl<'a> WheelFinder<'a> { - /// Initialize a new wheel finder. +impl<'a> DistributionFinder<'a> { + /// Initialize a new distribution finder. pub fn new(tags: &'a Tags, client: &'a RegistryClient) -> Self { Self { tags, @@ -53,23 +51,16 @@ impl<'a> WheelFinder<'a> { return Ok(Resolution::default()); } - // A channel to fetch package metadata (e.g., given `flask`, fetch all versions) and version - // metadata (e.g., given `flask==1.0.0`, fetch the metadata for that version). + // A channel to fetch package metadata (e.g., given `flask`, fetch all versions). let (package_sink, package_stream) = futures::channel::mpsc::unbounded(); // Initialize the package stream. let mut package_stream = package_stream .map(|request: Request| match request { - Request::Package(requirement) => Either::Left( - self.client - .simple(requirement.name.clone()) - .map_ok(move |metadata| Response::Package(requirement, metadata)), - ), - Request::Version(requirement, file) => Either::Right( - self.client - .file(file.clone()) - .map_ok(move |metadata| Response::Version(requirement, file, metadata)), - ), + Request::Package(requirement) => self + .client + .simple(requirement.name.clone()) + .map_ok(move |metadata| Response::Package(requirement, metadata)), }) .buffer_unordered(32) .ready_chunks(32); @@ -107,42 +98,17 @@ impl<'a> WheelFinder<'a> { match result { Response::Package(requirement, metadata) => { // Pick a version that satisfies the requirement. - let Some(file) = metadata.files.iter().rev().find(|file| { - // We only support wheels for now. - let Ok(name) = WheelFilename::from_str(file.filename.as_str()) else { - return false; - }; - - if !name.is_compatible(self.tags) { - return false; - } - - requirement.is_satisfied_by(&name.version) - }) else { + let Some(distribution) = self.select(&requirement, metadata.files) else { return Err(ResolveError::NotFound(requirement)); }; - package_sink.unbounded_send(Request::Version(requirement, file.clone()))?; - } - Response::Version(requirement, file, metadata) => { - debug!( - "Selecting: {}=={} ({})", - metadata.name, metadata.version, file.filename - ); - - let package = RemoteDistribution::from_registry( - metadata.name, - metadata.version, - file, - ); - if let Some(reporter) = self.reporter.as_ref() { - reporter.on_progress(&package); + reporter.on_progress(&distribution); } // Add to the resolved set. let normalized_name = requirement.name.clone(); - resolution.insert(normalized_name, package); + resolution.insert(normalized_name, distribution); } } } @@ -158,26 +124,52 @@ impl<'a> WheelFinder<'a> { Ok(Resolution::new(resolution)) } + + /// select a version that satisfies the requirement, preferring wheels to source distributions. + fn select(&self, requirement: &Requirement, files: Vec) -> Option { + let mut fallback = None; + for file in files.into_iter().rev() { + if let Ok(wheel) = WheelFilename::from_str(file.filename.as_str()) { + if !wheel.is_compatible(self.tags) { + continue; + } + if requirement.is_satisfied_by(&wheel.version) { + return Some(RemoteDistribution::from_registry( + wheel.distribution, + wheel.version, + file, + )); + } + } else if let Ok(sdist) = + SourceDistributionFilename::parse(file.filename.as_str(), &requirement.name) + { + if requirement.is_satisfied_by(&sdist.version) { + fallback = Some(RemoteDistribution::from_registry( + sdist.name, + sdist.version, + file, + )); + } + } + } + fallback + } } #[derive(Debug)] enum Request { /// A request to fetch the metadata for a package. Package(Requirement), - /// A request to fetch the metadata for a specific version of a package. - Version(Requirement, File), } #[derive(Debug)] enum Response { /// The returned metadata for a package. Package(Requirement, SimpleJson), - /// The returned metadata for a specific version of a package. - Version(Requirement, File, Metadata21), } pub trait Reporter: Send + Sync { - /// Callback to invoke when a package is resolved to a wheel. + /// Callback to invoke when a package is resolved to a specific distribution. fn on_progress(&self, wheel: &RemoteDistribution); /// Callback to invoke when the resolution is complete. diff --git a/crates/puffin-resolver/src/lib.rs b/crates/puffin-resolver/src/lib.rs index aa03847c5..f936b6614 100644 --- a/crates/puffin-resolver/src/lib.rs +++ b/crates/puffin-resolver/src/lib.rs @@ -1,20 +1,20 @@ pub use error::ResolveError; +pub use finder::{DistributionFinder, Reporter as FinderReporter}; pub use manifest::Manifest; pub use prerelease_mode::PreReleaseMode; pub use pubgrub::ResolutionFailureReporter; pub use resolution::Graph; pub use resolution_mode::ResolutionMode; pub use resolver::{Reporter as ResolverReporter, Resolver}; -pub use wheel_finder::{Reporter as WheelFinderReporter, WheelFinder}; mod candidate_selector; mod distribution; mod error; mod file; +mod finder; mod manifest; mod prerelease_mode; mod pubgrub; mod resolution; mod resolution_mode; mod resolver; -mod wheel_finder;