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

@@ -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)
+        let wheel_finder = puffin_resolver::DistributionFinder::new(&tags, &client)
-            .with_reporter(WheelFinderReporter::from(printer).with_length(remote.len() as u64));
+            .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" };

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<Printer> for WheelFinderReporter {
+impl From<Printer> 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<Printer> 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);

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")?;

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<Box<dyn Reporter>>,
 }
 
-impl<'a> WheelFinder<'a> {
+impl<'a> DistributionFinder<'a> {
-    /// Initialize a new wheel finder.
+    /// 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
+        // A channel to fetch package metadata (e.g., given `flask`, fetch all versions).
-        // metadata (e.g., given `flask==1.0.0`, fetch the metadata for that version).
         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(
+                Request::Package(requirement) => self
-                    self.client
+                    .client
-                        .simple(requirement.name.clone())
+                    .simple(requirement.name.clone())
-                        .map_ok(move |metadata| Response::Package(requirement, metadata)),
+                    .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)),
-                ),
             })
             .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| {
+                        let Some(distribution) = self.select(&requirement, metadata.files) else {
-                            // 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 {
                             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<File>) -> Option<RemoteDistribution> {
+        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.

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;