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

@@ -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::<Result<Vec<Requirement>>>()?;
+    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,

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<PathBuf>,
 
+    /// Constrain versions using the given constraints files.
+    #[clap(short, long)]
+    constraint: Vec<PathBuf>,
+
     /// Write the compiled requirements to the given `requirements.txt` file.
     #[clap(short, long)]
     output_file: Option<PathBuf>,
@@ -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<PathBuf>,
+
     /// 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::<Vec<_>>();
+            let constraints = args
+                .constraint
+                .into_iter()
+                .map(RequirementsSource::from)
+                .collect::<Vec<_>>();
             commands::pip_compile(
-                &sources,
+                &requirements,
+                &constraints,
                 args.output_file.as_deref(),
                 dirs.as_ref()
                     .map(ProjectDirs::cache_dir)

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<Range<PubGrubVersion>> {
+pub(crate) fn version_range(
+    specifiers: Option<&pep508_rs::VersionOrUrl>,
+) -> Result<Range<PubGrubVersion>> {
     let Some(specifiers) = specifiers else {
         return Ok(Range::any());
     };

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<Requirement>,
+    constraints: Vec<Requirement>,
     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<Requirement>,
+        constraints: Vec<Requirement>,
         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

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