Add support for lowest and lowest-direct resolution modes (#160)

Borrows terminology from pnpm by introducing three resolution modes:

- "Highest": always choose the highest compliant version (default).
- "Lowest": always choose the lowest compliant version.
- "LowestDirect": choose the lowest compliant version of direct
dependencies, and the highest compliant version of any transitive
dependencies. (This makes a bit more sense than "lowest".)

Closes https://github.com/astral-sh/puffin/issues/142.

Cargo.lock

@@ -2173,11 +2173,13 @@ version = "0.0.1"
 dependencies = [
  "anyhow",
  "bitflags 2.4.1",
+ "clap",
  "colored",
  "distribution-filename",
  "futures",
  "fxhash",
  "insta",
+ "itertools",
  "once_cell",
  "pep440_rs 0.3.12",
  "pep508_rs",

crates/puffin-cli/Cargo.toml

@@ -19,7 +19,7 @@ puffin-client = { path = "../puffin-client" }
 puffin-installer = { path = "../puffin-installer" }
 puffin-interpreter = { path = "../puffin-interpreter" }
 puffin-package = { path = "../puffin-package" }
-puffin-resolver = { path = "../puffin-resolver" }
+puffin-resolver = { path = "../puffin-resolver", features = ["clap"] }
 puffin-workspace = { path = "../puffin-workspace" }
 
 anyhow = { workspace = true }

crates/puffin-cli/src/commands/pip_compile.rs

@@ -15,6 +15,7 @@ use platform_host::Platform;
 use platform_tags::Tags;
 use puffin_client::PypiClientBuilder;
 use puffin_interpreter::PythonExecutable;
+use puffin_resolver::ResolutionMode;
 
 use crate::commands::{elapsed, ExitStatus};
 use crate::printer::Printer;
@@ -27,6 +28,7 @@ pub(crate) async fn pip_compile(
     requirements: &[RequirementsSource],
     constraints: &[RequirementsSource],
     output_file: Option<&Path>,
+    mode: ResolutionMode,
     cache: Option<&Path>,
     mut printer: Printer,
 ) -> Result<ExitStatus> {
@@ -63,7 +65,7 @@ pub(crate) async fn pip_compile(
 
     // Resolve the dependencies.
     let resolver =
-        puffin_resolver::Resolver::new(requirements, constraints, markers, &tags, &client);
+        puffin_resolver::Resolver::new(requirements, constraints, mode, markers, &tags, &client);
     let resolution = match resolver.resolve().await {
         Err(puffin_resolver::ResolveError::PubGrub(pubgrub::error::PubGrubError::NoSolution(
             mut derivation_tree,

crates/puffin-cli/src/main.rs

@@ -4,6 +4,7 @@ use std::process::ExitCode;
 use clap::{Args, Parser, Subcommand};
 use colored::Colorize;
 use directories::ProjectDirs;
+use puffin_resolver::ResolutionMode;
 
 use crate::commands::ExitStatus;
 use crate::requirements::RequirementsSource;
@@ -63,6 +64,9 @@ struct PipCompileArgs {
     #[clap(short, long)]
     constraint: Vec<PathBuf>,
 
+    #[clap(long, value_enum)]
+    resolution: Option<ResolutionMode>,
+
     /// Write the compiled requirements to the given `requirements.txt` file.
     #[clap(short, long)]
     output_file: Option<PathBuf>,
@@ -146,6 +150,7 @@ async fn main() -> ExitCode {
                 &requirements,
                 &constraints,
                 args.output_file.as_deref(),
+                args.resolution.unwrap_or_default(),
                 dirs.as_ref()
                     .map(ProjectDirs::cache_dir)
                     .filter(|_| !cli.no_cache),

crates/puffin-resolver/Cargo.toml

@@ -21,9 +21,11 @@ distribution-filename = { path = "../distribution-filename" }
 
 anyhow = { workspace = true }
 bitflags = { workspace = true }
+clap = { workspace = true, features = ["derive"], optional = true }
 colored = { workspace = true }
 futures = { workspace = true }
 fxhash = { workspace = true }
+itertools = { workspace = true }
 once_cell = { workspace = true }
 petgraph = { workspace = true }
 thiserror = { workspace = true }

crates/puffin-resolver/src/lib.rs

@@ -1,9 +1,11 @@
 pub use error::ResolveError;
+pub use mode::ResolutionMode;
 pub use resolution::PinnedPackage;
 pub use resolver::Resolver;
 pub use wheel_finder::{Reporter, WheelFinder};
 
 mod error;
+mod mode;
 mod pubgrub;
 mod resolution;
 mod resolver;

crates/puffin-resolver/src/mode.rs

@@ -0,0 +1,67 @@
+use fxhash::FxHashSet;
+use itertools::Either;
+
+use pep508_rs::Requirement;
+use puffin_client::File;
+use puffin_package::package_name::PackageName;
+
+#[derive(Debug, Default, Clone, Copy, PartialEq, Eq)]
+#[cfg_attr(feature = "clap", derive(clap::ValueEnum))]
+pub enum ResolutionMode {
+    /// Resolve the highest compatible version of each package.
+    #[default]
+    Highest,
+    /// Resolve the lowest compatible version of each package.
+    Lowest,
+    /// Resolve the lowest compatible version of any direct dependencies, and the highest
+    /// compatible version of any transitive dependencies.
+    LowestDirect,
+}
+
+#[derive(Debug, Clone)]
+pub(crate) enum CandidateSelector {
+    /// Resolve the highest compatible version of each package.
+    Highest,
+    /// Resolve the lowest compatible version of each package.
+    Lowest,
+    /// Resolve the lowest compatible version of any direct dependencies, and the highest
+    /// compatible version of any transitive dependencies.
+    LowestDirect(FxHashSet<PackageName>),
+}
+
+impl CandidateSelector {
+    /// Return a candidate selector for the given resolution mode.
+    pub(crate) fn from_mode(mode: ResolutionMode, direct_dependencies: &[Requirement]) -> Self {
+        match mode {
+            ResolutionMode::Highest => Self::Highest,
+            ResolutionMode::Lowest => Self::Lowest,
+            ResolutionMode::LowestDirect => Self::LowestDirect(
+                direct_dependencies
+                    .iter()
+                    .map(|requirement| PackageName::normalize(&requirement.name))
+                    .collect(),
+            ),
+        }
+    }
+}
+
+impl CandidateSelector {
+    /// Return an iterator over the candidates for the given package name.
+    pub(crate) fn iter_candidates<'a>(
+        &self,
+        package_name: &PackageName,
+        candidates: &'a [File],
+    ) -> impl Iterator<Item = &'a File> {
+        match self {
+            CandidateSelector::Highest => Either::Left(candidates.iter().rev()),
+            CandidateSelector::Lowest => Either::Right(candidates.iter()),
+            CandidateSelector::LowestDirect(direct_dependencies) => {
+                if direct_dependencies.contains(package_name) {
+                    Either::Right(candidates.iter())
+                } else {
+                    Either::Left(candidates.iter().rev())
+                }
+            }
+        }
+    }
+}

crates/puffin-resolver/src/resolver.rs

@@ -27,6 +27,7 @@ use puffin_package::metadata::Metadata21;
 use puffin_package::package_name::PackageName;
 
 use crate::error::ResolveError;
+use crate::mode::{CandidateSelector, ResolutionMode};
 use crate::pubgrub::package::PubGrubPackage;
 use crate::pubgrub::version::{PubGrubVersion, MIN_VERSION};
 use crate::pubgrub::{iter_requirements, version_range};
@@ -38,6 +39,7 @@ pub struct Resolver<'a> {
     markers: &'a MarkerEnvironment,
     tags: &'a Tags,
     client: &'a PypiClient,
+    selector: CandidateSelector,
     cache: Arc<SolverCache>,
 }
 
@@ -46,17 +48,19 @@ impl<'a> Resolver<'a> {
     pub fn new(
         requirements: Vec<Requirement>,
         constraints: Vec<Requirement>,
+        mode: ResolutionMode,
         markers: &'a MarkerEnvironment,
         tags: &'a Tags,
         client: &'a PypiClient,
     ) -> Self {
         Self {
+            selector: CandidateSelector::from_mode(mode, &requirements),
+            cache: Arc::new(SolverCache::default()),
             requirements,
             constraints,
             markers,
             tags,
             client,
-            cache: Arc::new(SolverCache::default()),
         }
     }
 
@@ -263,24 +267,28 @@ impl<'a> Resolver<'a> {
             let simple_json = entry.value();
 
             // Select the latest compatible version.
-            let Some(file) = simple_json.files.iter().rev().find(|file| {
+            let Some(file) = self
-                let Ok(name) = WheelFilename::from_str(file.filename.as_str()) else {
+                .selector
-                    return false;
+                .iter_candidates(package_name, &simple_json.files)
-                };
+                .find(|file| {
+                    let Ok(name) = WheelFilename::from_str(file.filename.as_str()) else {
+                        return false;
+                    };
 
-                if !name.is_compatible(self.tags) {
+                    if !name.is_compatible(self.tags) {
-                    return false;
+                        return false;
-                }
+                    }
 
-                if !range
+                    if !range
-                    .borrow()
+                        .borrow()
-                    .contains(&PubGrubVersion::from(name.version.clone()))
+                        .contains(&PubGrubVersion::from(name.version.clone()))
-                {
+                    {
-                    return false;
+                        return false;
-                };
+                    };
 
-                true
+                    true
-            }) else {
+                })
+            else {
                 // Short circuit: we couldn't find _any_ compatible versions for a package.
                 let (package, _range) = potential_packages.swap_remove(index);
                 return Ok((package, None));
@@ -313,28 +321,32 @@ impl<'a> Resolver<'a> {
                 );
 
                 // Find a compatible version.
-                let Some(wheel) = simple_json.files.iter().rev().find_map(|file| {
+                let Some(wheel) = self
-                    let Ok(name) = WheelFilename::from_str(file.filename.as_str()) else {
+                    .selector
-                        return None;
+                    .iter_candidates(package_name, &simple_json.files)
-                    };
+                    .find_map(|file| {
+                        let Ok(name) = WheelFilename::from_str(file.filename.as_str()) else {
+                            return None;
+                        };
 
-                    if !name.is_compatible(self.tags) {
+                        if !name.is_compatible(self.tags) {
-                        return None;
+                            return None;
-                    }
+                        }
 
-                    if !range
+                        if !range
-                        .borrow()
+                            .borrow()
-                        .contains(&PubGrubVersion::from(name.version.clone()))
+                            .contains(&PubGrubVersion::from(name.version.clone()))
-                    {
+                        {
-                        return None;
+                            return None;
-                    };
+                        };
 
-                    Some(Wheel {
+                        Some(Wheel {
-                        file: file.clone(),
+                            file: file.clone(),
-                        name: package_name.clone(),
+                            name: package_name.clone(),
-                        version: name.version.clone(),
+                            version: name.version.clone(),
+                        })
                     })
-                }) else {
+                else {
                     // Short circuit: we couldn't find _any_ compatible versions for a package.
                     return Ok((package, None));
                 };

crates/puffin-resolver/tests/resolver.rs

@@ -10,7 +10,7 @@ use pep508_rs::{MarkerEnvironment, Requirement, StringVersion};
 use platform_host::{Arch, Os, Platform};
 use platform_tags::Tags;
 use puffin_client::PypiClientBuilder;
-use puffin_resolver::Resolver;
+use puffin_resolver::{ResolutionMode, Resolver};
 
 #[tokio::test]
 async fn pylint() -> Result<()> {
@@ -20,7 +20,14 @@ async fn pylint() -> Result<()> {
 
     let requirements = vec![Requirement::from_str("pylint==2.3.0").unwrap()];
     let constraints = vec![];
-    let resolver = Resolver::new(requirements, constraints, &MARKERS_311, &TAGS_311, &client);
+    let resolver = Resolver::new(
+        requirements,
+        constraints,
+        ResolutionMode::default(),
+        &MARKERS_311,
+        &TAGS_311,
+        &client,
+    );
     let resolution = resolver.resolve().await?;
 
     insta::assert_display_snapshot!(resolution);
@@ -36,7 +43,14 @@ async fn black() -> Result<()> {
 
     let requirements = vec![Requirement::from_str("black<=23.9.1").unwrap()];
     let constraints = vec![];
-    let resolver = Resolver::new(requirements, constraints, &MARKERS_311, &TAGS_311, &client);
+    let resolver = Resolver::new(
+        requirements,
+        constraints,
+        ResolutionMode::default(),
+        &MARKERS_311,
+        &TAGS_311,
+        &client,
+    );
     let resolution = resolver.resolve().await?;
 
     insta::assert_display_snapshot!(resolution);
@@ -52,7 +66,14 @@ async fn black_colorama() -> Result<()> {
 
     let requirements = vec![Requirement::from_str("black[colorama]<=23.9.1").unwrap()];
     let constraints = vec![];
-    let resolver = Resolver::new(requirements, constraints, &MARKERS_311, &TAGS_311, &client);
+    let resolver = Resolver::new(
+        requirements,
+        constraints,
+        ResolutionMode::default(),
+        &MARKERS_311,
+        &TAGS_311,
+        &client,
+    );
     let resolution = resolver.resolve().await?;
 
     insta::assert_display_snapshot!(resolution);
@@ -68,7 +89,14 @@ async fn black_python_310() -> Result<()> {
 
     let requirements = vec![Requirement::from_str("black<=23.9.1").unwrap()];
     let constraints = vec![];
-    let resolver = Resolver::new(requirements, constraints, &MARKERS_310, &TAGS_310, &client);
+    let resolver = Resolver::new(
+        requirements,
+        constraints,
+        ResolutionMode::default(),
+        &MARKERS_310,
+        &TAGS_310,
+        &client,
+    );
     let resolution = resolver.resolve().await?;
 
     insta::assert_display_snapshot!(resolution);
@@ -86,7 +114,14 @@ async fn black_mypy_extensions() -> Result<()> {
 
     let requirements = vec![Requirement::from_str("black<=23.9.1").unwrap()];
     let constraints = vec![Requirement::from_str("mypy-extensions<1").unwrap()];
-    let resolver = Resolver::new(requirements, constraints, &MARKERS_311, &TAGS_311, &client);
+    let resolver = Resolver::new(
+        requirements,
+        constraints,
+        ResolutionMode::default(),
+        &MARKERS_311,
+        &TAGS_311,
+        &client,
+    );
     let resolution = resolver.resolve().await?;
 
     insta::assert_display_snapshot!(resolution);
@@ -104,7 +139,14 @@ async fn black_mypy_extensions_extra() -> Result<()> {
 
     let requirements = vec![Requirement::from_str("black<=23.9.1").unwrap()];
     let constraints = vec![Requirement::from_str("mypy-extensions[extra]<1").unwrap()];
-    let resolver = Resolver::new(requirements, constraints, &MARKERS_311, &TAGS_311, &client);
+    let resolver = Resolver::new(
+        requirements,
+        constraints,
+        ResolutionMode::default(),
+        &MARKERS_311,
+        &TAGS_311,
+        &client,
+    );
     let resolution = resolver.resolve().await?;
 
     insta::assert_display_snapshot!(resolution);
@@ -122,7 +164,60 @@ async fn black_flake8() -> Result<()> {
 
     let requirements = vec![Requirement::from_str("black<=23.9.1").unwrap()];
     let constraints = vec![Requirement::from_str("flake8<1").unwrap()];
-    let resolver = Resolver::new(requirements, constraints, &MARKERS_311, &TAGS_311, &client);
+    let resolver = Resolver::new(
+        requirements,
+        constraints,
+        ResolutionMode::default(),
+        &MARKERS_311,
+        &TAGS_311,
+        &client,
+    );
+    let resolution = resolver.resolve().await?;
+
+    insta::assert_display_snapshot!(resolution);
+
+    Ok(())
+}
+
+#[tokio::test]
+async fn black_lowest() -> Result<()> {
+    colored::control::set_override(false);
+
+    let client = PypiClientBuilder::default().build();
+
+    let requirements = vec![Requirement::from_str("black>21").unwrap()];
+    let constraints = vec![];
+    let resolver = Resolver::new(
+        requirements,
+        constraints,
+        ResolutionMode::Lowest,
+        &MARKERS_311,
+        &TAGS_311,
+        &client,
+    );
+    let resolution = resolver.resolve().await?;
+
+    insta::assert_display_snapshot!(resolution);
+
+    Ok(())
+}
+
+#[tokio::test]
+async fn black_lowest_direct() -> Result<()> {
+    colored::control::set_override(false);
+
+    let client = PypiClientBuilder::default().build();
+
+    let requirements = vec![Requirement::from_str("black>21").unwrap()];
+    let constraints = vec![];
+    let resolver = Resolver::new(
+        requirements,
+        constraints,
+        ResolutionMode::LowestDirect,
+        &MARKERS_311,
+        &TAGS_311,
+        &client,
+    );
     let resolution = resolver.resolve().await?;
 
     insta::assert_display_snapshot!(resolution);

crates/puffin-resolver/tests/snapshots/resolver__black_lowest.snap

@@ -0,0 +1,18 @@
+---
+source: crates/puffin-resolver/tests/resolver.rs
+expression: resolution
+---
+appdirs==1.4.0
+    # via black
+black==21.4b0
+click==7.1.2
+    # via black
+mypy-extensions==0.4.3
+    # via black
+pathspec==0.7.0
+    # via black
+regex==2022.9.11
+    # via black
+toml==0.10.1
+    # via black
+

crates/puffin-resolver/tests/snapshots/resolver__black_lowest_direct.snap

@@ -0,0 +1,18 @@
+---
+source: crates/puffin-resolver/tests/resolver.rs
+expression: resolution
+---
+appdirs==1.4.4
+    # via black
+black==21.4b0
+click==8.1.7
+    # via black
+mypy-extensions==1.0.0
+    # via black
+pathspec==0.11.2
+    # via black
+regex==2023.10.3
+    # via black
+toml==0.10.2
+    # via black
+