diff --git a/Cargo.lock b/Cargo.lock index 024944a4b..130e83e7d 100644 --- a/Cargo.lock +++ b/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", diff --git a/crates/puffin-cli/Cargo.toml b/crates/puffin-cli/Cargo.toml index 7d154427e..e13fc2050 100644 --- a/crates/puffin-cli/Cargo.toml +++ b/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 } diff --git a/crates/puffin-cli/src/commands/pip_compile.rs b/crates/puffin-cli/src/commands/pip_compile.rs index b7ecc675c..1808c8d03 100644 --- a/crates/puffin-cli/src/commands/pip_compile.rs +++ b/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 { @@ -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, diff --git a/crates/puffin-cli/src/main.rs b/crates/puffin-cli/src/main.rs index 0bbced5c7..37a526a62 100644 --- a/crates/puffin-cli/src/main.rs +++ b/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, + #[clap(long, value_enum)] + resolution: Option, + /// Write the compiled requirements to the given `requirements.txt` file. #[clap(short, long)] output_file: Option, @@ -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), diff --git a/crates/puffin-resolver/Cargo.toml b/crates/puffin-resolver/Cargo.toml index f9a75fa9b..25443a650 100644 --- a/crates/puffin-resolver/Cargo.toml +++ b/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 } diff --git a/crates/puffin-resolver/src/lib.rs b/crates/puffin-resolver/src/lib.rs index c5fc93bbe..df1c19a46 100644 --- a/crates/puffin-resolver/src/lib.rs +++ b/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; diff --git a/crates/puffin-resolver/src/mode.rs b/crates/puffin-resolver/src/mode.rs new file mode 100644 index 000000000..064960a79 --- /dev/null +++ b/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), +} + +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 { + 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()) + } + } + } + } +} diff --git a/crates/puffin-resolver/src/resolver.rs b/crates/puffin-resolver/src/resolver.rs index c537bdb11..0cb89c079 100644 --- a/crates/puffin-resolver/src/resolver.rs +++ b/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, } @@ -46,17 +48,19 @@ impl<'a> Resolver<'a> { pub fn new( requirements: Vec, constraints: Vec, + 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 Ok(name) = WheelFilename::from_str(file.filename.as_str()) else { - return false; - }; + let Some(file) = self + .selector + .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) { - return false; - } + if !name.is_compatible(self.tags) { + return false; + } - if !range - .borrow() - .contains(&PubGrubVersion::from(name.version.clone())) - { - return false; - }; + if !range + .borrow() + .contains(&PubGrubVersion::from(name.version.clone())) + { + return false; + }; - true - }) else { + true + }) + 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 Ok(name) = WheelFilename::from_str(file.filename.as_str()) else { - return None; - }; + let Some(wheel) = self + .selector + .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) { - return None; - } + if !name.is_compatible(self.tags) { + return None; + } - if !range - .borrow() - .contains(&PubGrubVersion::from(name.version.clone())) - { - return None; - }; + if !range + .borrow() + .contains(&PubGrubVersion::from(name.version.clone())) + { + return None; + }; - Some(Wheel { - file: file.clone(), - name: package_name.clone(), - version: name.version.clone(), + Some(Wheel { + file: file.clone(), + name: package_name.clone(), + version: name.version.clone(), + }) }) - }) else { + else { // Short circuit: we couldn't find _any_ compatible versions for a package. return Ok((package, None)); }; diff --git a/crates/puffin-resolver/tests/resolver.rs b/crates/puffin-resolver/tests/resolver.rs index 8e7a7de18..da5e7c737 100644 --- a/crates/puffin-resolver/tests/resolver.rs +++ b/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); diff --git a/crates/puffin-resolver/tests/snapshots/resolver__black_lowest.snap b/crates/puffin-resolver/tests/snapshots/resolver__black_lowest.snap new file mode 100644 index 000000000..a8485dd3b --- /dev/null +++ b/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 + diff --git a/crates/puffin-resolver/tests/snapshots/resolver__black_lowest_direct.snap b/crates/puffin-resolver/tests/snapshots/resolver__black_lowest_direct.snap new file mode 100644 index 000000000..b9c96047d --- /dev/null +++ b/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 +