From 0b60804db6e64021a055db34c0491e7eb129e5d4 Mon Sep 17 00:00:00 2001 From: Charlie Marsh Date: Thu, 19 Oct 2023 20:24:05 -0400 Subject: [PATCH] Add support for constraints during pip-compile resolution (#144) Closes https://github.com/astral-sh/puffin/issues/130. --- crates/puffin-cli/src/commands/pip_compile.rs | 13 ++- crates/puffin-cli/src/main.rs | 21 +++- crates/puffin-resolver/src/pubgrub/mod.rs | 4 +- crates/puffin-resolver/src/resolver.rs | 30 +++++- crates/puffin-resolver/tests/resolver.rs | 96 ++++++++++++++++++- 5 files changed, 149 insertions(+), 15 deletions(-) diff --git a/crates/puffin-cli/src/commands/pip_compile.rs b/crates/puffin-cli/src/commands/pip_compile.rs index 632d3de00..a871105e6 100644 --- a/crates/puffin-cli/src/commands/pip_compile.rs +++ b/crates/puffin-cli/src/commands/pip_compile.rs @@ -21,7 +21,8 @@ use crate::requirements::RequirementsSource; /// Resolve a set of requirements into a set of pinned versions. pub(crate) async fn pip_compile( - sources: &[RequirementsSource], + requirements: &[RequirementsSource], + constraints: &[RequirementsSource], output_file: Option<&Path>, cache: Option<&Path>, mut printer: Printer, @@ -29,7 +30,12 @@ pub(crate) async fn pip_compile( let start = std::time::Instant::now(); // Read all requirements from the provided sources. - let requirements = sources + let requirements = requirements + .iter() + .map(RequirementsSource::requirements) + .flatten_ok() + .collect::>>()?; + let constraints = constraints .iter() .map(RequirementsSource::requirements) .flatten_ok() @@ -53,7 +59,8 @@ pub(crate) async fn pip_compile( let client = PypiClientBuilder::default().cache(cache).build(); // Resolve the dependencies. - let resolver = puffin_resolver::Resolver::new(requirements, markers, &tags, &client); + let resolver = + puffin_resolver::Resolver::new(requirements, constraints, 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 6e4e7ec3f..a5d6faf26 100644 --- a/crates/puffin-cli/src/main.rs +++ b/crates/puffin-cli/src/main.rs @@ -39,12 +39,12 @@ enum Commands { PipCompile(PipCompileArgs), /// Sync dependencies from a `requirements.txt` file. PipSync(PipSyncArgs), + /// Uninstall packages from the current environment. + PipUninstall(PipUninstallArgs), /// Clear the cache. Clean, /// Enumerate the installed packages in the current environment. Freeze, - /// Uninstall packages from the current environment. - PipUninstall(PipUninstallArgs), /// Create a virtual environment. Venv(VenvArgs), /// Add a dependency to the workspace. @@ -59,6 +59,10 @@ struct PipCompileArgs { #[clap(required(true))] src_file: Vec, + /// Constrain versions using the given constraints files. + #[clap(short, long)] + constraint: Vec, + /// Write the compiled requirements to the given `requirements.txt` file. #[clap(short, long)] output_file: Option, @@ -85,11 +89,12 @@ struct PipUninstallArgs { #[derive(Args)] struct VenvArgs { - /// The python interpreter to use for the virtual environment + /// The Python interpreter to use for the virtual environment. // Short `-p` to match `virtualenv` // TODO(konstin): Support e.g. `-p 3.10` #[clap(short, long)] python: Option, + /// The path to the virtual environment to create. name: PathBuf, } @@ -127,13 +132,19 @@ async fn main() -> ExitCode { let result = match cli.command { Commands::PipCompile(args) => { let dirs = ProjectDirs::from("", "", "puffin"); - let sources = args + let requirements = args .src_file .into_iter() .map(RequirementsSource::from) .collect::>(); + let constraints = args + .constraint + .into_iter() + .map(RequirementsSource::from) + .collect::>(); commands::pip_compile( - &sources, + &requirements, + &constraints, args.output_file.as_deref(), dirs.as_ref() .map(ProjectDirs::cache_dir) diff --git a/crates/puffin-resolver/src/pubgrub/mod.rs b/crates/puffin-resolver/src/pubgrub/mod.rs index e0dabb02f..3b7321a00 100644 --- a/crates/puffin-resolver/src/pubgrub/mod.rs +++ b/crates/puffin-resolver/src/pubgrub/mod.rs @@ -54,7 +54,9 @@ pub(crate) fn iter_requirements<'a>( } /// Convert a PEP 508 specifier to a `PubGrub` range. -fn version_range(specifiers: Option<&pep508_rs::VersionOrUrl>) -> Result> { +pub(crate) fn version_range( + specifiers: Option<&pep508_rs::VersionOrUrl>, +) -> Result> { let Some(specifiers) = specifiers else { return Ok(Range::any()); }; diff --git a/crates/puffin-resolver/src/resolver.rs b/crates/puffin-resolver/src/resolver.rs index 733b84b04..72b352d20 100644 --- a/crates/puffin-resolver/src/resolver.rs +++ b/crates/puffin-resolver/src/resolver.rs @@ -27,13 +27,14 @@ use puffin_package::package_name::PackageName; use wheel_filename::WheelFilename; use crate::error::ResolveError; -use crate::pubgrub::iter_requirements; use crate::pubgrub::package::PubGrubPackage; use crate::pubgrub::version::{PubGrubVersion, MIN_VERSION}; +use crate::pubgrub::{iter_requirements, version_range}; use crate::resolution::{PinnedPackage, Resolution}; pub struct Resolver<'a> { requirements: Vec, + constraints: Vec, markers: &'a MarkerEnvironment, tags: &'a Tags, client: &'a PypiClient, @@ -44,12 +45,14 @@ impl<'a> Resolver<'a> { /// Initialize a new resolver. pub fn new( requirements: Vec, + constraints: Vec, markers: &'a MarkerEnvironment, tags: &'a Tags, client: &'a PypiClient, ) -> Self { Self { requirements, + constraints, markers, tags, client, @@ -376,6 +379,8 @@ impl<'a> Resolver<'a> { match package { PubGrubPackage::Root => { let mut constraints = DependencyConstraints::default(); + + // Add the root requirements. for (package, version) in iter_requirements(self.requirements.iter(), None, self.markers) { @@ -388,6 +393,18 @@ impl<'a> Resolver<'a> { } } } + + // If any requirements were further constrained by the user, add those constraints. + for constraint in &self.constraints { + let package = + PubGrubPackage::Package(PackageName::normalize(&constraint.name), None); + if let Some(range) = constraints.get_mut(&package) { + *range = range.intersection( + &version_range(constraint.version_or_url.as_ref()).unwrap(), + ); + } + } + Ok(Dependencies::Known(constraints)) } PubGrubPackage::Package(package_name, extra) => { @@ -427,6 +444,17 @@ impl<'a> Resolver<'a> { } } + // If any packages were further constrained by the user, add those constraints. + for constraint in &self.constraints { + let package = + PubGrubPackage::Package(PackageName::normalize(&constraint.name), None); + if let Some(range) = constraints.get_mut(&package) { + *range = range.intersection( + &version_range(constraint.version_or_url.as_ref()).unwrap(), + ); + } + } + if let Some(extra) = extra { if !metadata .provides_extras diff --git a/crates/puffin-resolver/tests/resolver.rs b/crates/puffin-resolver/tests/resolver.rs index d15261f9b..df7738867 100644 --- a/crates/puffin-resolver/tests/resolver.rs +++ b/crates/puffin-resolver/tests/resolver.rs @@ -17,7 +17,8 @@ async fn pylint() -> Result<()> { let client = PypiClientBuilder::default().build(); let requirements = vec![Requirement::from_str("pylint==2.3.0").unwrap()]; - let resolver = Resolver::new(requirements, &MARKERS_311, &TAGS_311, &client); + let constraints = vec![]; + let resolver = Resolver::new(requirements, constraints, &MARKERS_311, &TAGS_311, &client); let resolution = resolver.resolve().await?; assert_eq!( @@ -39,7 +40,8 @@ async fn black() -> Result<()> { let client = PypiClientBuilder::default().build(); let requirements = vec![Requirement::from_str("black<=23.9.1").unwrap()]; - let resolver = Resolver::new(requirements, &MARKERS_311, &TAGS_311, &client); + let constraints = vec![]; + let resolver = Resolver::new(requirements, constraints, &MARKERS_311, &TAGS_311, &client); let resolution = resolver.resolve().await?; assert_eq!( @@ -63,7 +65,8 @@ async fn black_colorama() -> Result<()> { let client = PypiClientBuilder::default().build(); let requirements = vec![Requirement::from_str("black[colorama]<=23.9.1").unwrap()]; - let resolver = Resolver::new(requirements, &MARKERS_311, &TAGS_311, &client); + let constraints = vec![]; + let resolver = Resolver::new(requirements, constraints, &MARKERS_311, &TAGS_311, &client); let resolution = resolver.resolve().await?; assert_eq!( @@ -88,7 +91,8 @@ async fn black_python_310() -> Result<()> { let client = PypiClientBuilder::default().build(); let requirements = vec![Requirement::from_str("black<=23.9.1").unwrap()]; - let resolver = Resolver::new(requirements, &MARKERS_310, &TAGS_310, &client); + let constraints = vec![]; + let resolver = Resolver::new(requirements, constraints, &MARKERS_310, &TAGS_310, &client); let resolution = resolver.resolve().await?; assert_eq!( @@ -109,12 +113,94 @@ async fn black_python_310() -> Result<()> { Ok(()) } +/// Resolve `black` with a constraint on `mypy-extensions`, to ensure that constraints are +/// respected. +#[tokio::test] +async fn black_mypy_extensions() -> Result<()> { + let client = PypiClientBuilder::default().build(); + + 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 resolution = resolver.resolve().await?; + + assert_eq!( + format!("{resolution}"), + [ + "black==23.9.1", + "click==8.1.7", + "mypy-extensions==0.4.3", + "packaging==23.2", + "pathspec==0.11.2", + "platformdirs==3.11.0" + ] + .join("\n") + ); + + Ok(()) +} + +/// Resolve `black` with a constraint on `mypy-extensions[extra]`, to ensure that extras are +/// ignored when resolving constraints. +#[tokio::test] +async fn black_mypy_extensions_extra() -> Result<()> { + let client = PypiClientBuilder::default().build(); + + 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 resolution = resolver.resolve().await?; + + assert_eq!( + format!("{resolution}"), + [ + "black==23.9.1", + "click==8.1.7", + "mypy-extensions==0.4.3", + "packaging==23.2", + "pathspec==0.11.2", + "platformdirs==3.11.0" + ] + .join("\n") + ); + + Ok(()) +} + +/// Resolve `black` with a redundant constraint on `flake8`, to ensure that constraints don't +/// introduce new dependencies. +#[tokio::test] +async fn black_flake8() -> Result<()> { + let client = PypiClientBuilder::default().build(); + + 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 resolution = resolver.resolve().await?; + + assert_eq!( + format!("{resolution}"), + [ + "black==23.9.1", + "click==8.1.7", + "mypy-extensions==1.0.0", + "packaging==23.2", + "pathspec==0.11.2", + "platformdirs==3.11.0" + ] + .join("\n") + ); + + Ok(()) +} + #[tokio::test] async fn htmldate() -> Result<()> { let client = PypiClientBuilder::default().build(); let requirements = vec![Requirement::from_str("htmldate<=1.5.0").unwrap()]; - let resolver = Resolver::new(requirements, &MARKERS_311, &TAGS_311, &client); + let constraints = vec![]; + let resolver = Resolver::new(requirements, constraints, &MARKERS_311, &TAGS_311, &client); let resolution = resolver.resolve().await?; // Resolves to `htmldate==1.4.3` (rather than `htmldate==1.5.2`) because `htmldate==1.5.2` has