Enable package overrides in pip-compile (#631)

## Summary

This PR enables overrides to be passed to `pip-compile` and
`pip-install` via a new `--overrides` flag.

When overrides are provided, we effectively replace any requirements
that are overridden with the overridden versions. This is applied at all
depths of the tree.

The merge semantics are such that we replace _all_ requirements of a
package with _all_ requirements from the overrides files. So, for
example, if a package declares:

```
foo >= 1.0; python_version < '3.11'
foo < 1.0; python_version >= '3.11'
```

And the user provides an override like:
```
foo >= 2.0
```

Then _both_ of the `foo` requirements in the package will be replaced
with the override.

If instead, the user provided an override like:
```
foo >= 2.0; python_version < '3.11'
foo < 3.0; python_version >= '3.11'
```

Then we'd replace _both_ of the original `foo` requirements with both of
these overrides. (In technical terms, for each package in the
requirements file, we flat-map over its overrides.)

Closes #511.

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

@@ -36,6 +36,7 @@ const VERSION: &str = env!("CARGO_PKG_VERSION");
 pub(crate) async fn pip_compile(
     requirements: &[RequirementsSource],
     constraints: &[RequirementsSource],
+    overrides: &[RequirementsSource],
     extras: ExtrasSpecification<'_>,
     output_file: Option<&Path>,
     resolution_mode: ResolutionMode,
@@ -76,8 +77,9 @@ pub(crate) async fn pip_compile(
         project,
         requirements,
         constraints,
+        overrides,
         extras: used_extras,
-    } = RequirementsSpecification::from_sources(requirements, constraints, &extras)?;
+    } = RequirementsSpecification::from_sources(requirements, constraints, overrides, &extras)?;
 
     // Check that all provided extras are used
     if let ExtrasSpecification::Some(extras) = extras {
@@ -108,7 +110,7 @@ pub(crate) async fn pip_compile(
         .unwrap_or_default();
 
     // Create a manifest of the requirements.
-    let manifest = Manifest::new(requirements, constraints, preferences, project);
+    let manifest = Manifest::new(requirements, constraints, overrides, preferences, project);
     let options = ResolutionOptions::new(resolution_mode, prerelease_mode, exclude_newer);
 
     // Detect the current Python interpreter.

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

@@ -33,6 +33,7 @@ use crate::requirements::{ExtrasSpecification, RequirementsSource, RequirementsS
 pub(crate) async fn pip_install(
     requirements: &[RequirementsSource],
     constraints: &[RequirementsSource],
+    overrides: &[RequirementsSource],
     extras: &ExtrasSpecification<'_>,
     resolution_mode: ResolutionMode,
     prerelease_mode: PreReleaseMode,
@@ -57,7 +58,7 @@ pub(crate) async fn pip_install(
     let start = std::time::Instant::now();
 
     // Determine the requirements.
-    let spec = specification(requirements, constraints, extras)?;
+    let spec = specification(requirements, constraints, overrides, extras)?;
 
     // Detect the current Python interpreter.
     let platform = Platform::current()?;
@@ -123,6 +124,7 @@ pub(crate) async fn pip_install(
 fn specification(
     requirements: &[RequirementsSource],
     constraints: &[RequirementsSource],
+    overrides: &[RequirementsSource],
     extras: &ExtrasSpecification<'_>,
 ) -> Result<RequirementsSpecification> {
     // If the user requests `extras` but does not provide a pyproject toml source
@@ -137,7 +139,8 @@ fn specification(
     }
 
     // Read all requirements from the provided sources.
-    let spec = RequirementsSpecification::from_sources(requirements, constraints, extras)?;
+    let spec =
+        RequirementsSpecification::from_sources(requirements, constraints, overrides, extras)?;
 
     // Check that all provided extras are used
     if let ExtrasSpecification::Some(extras) = extras {
@@ -185,6 +188,7 @@ async fn resolve(
         project,
         requirements,
         constraints,
+        overrides,
         extras: _,
     } = spec;
 
@@ -200,7 +204,7 @@ async fn resolve(
             .collect(),
     };
 
-    let manifest = Manifest::new(requirements, constraints, preferences, project);
+    let manifest = Manifest::new(requirements, constraints, overrides, preferences, project);
     let options = ResolutionOptions::new(resolution_mode, prerelease_mode, exclude_newer);
 
     debug!(

crates/puffin-cli/src/main.rs

@@ -107,10 +107,28 @@ struct PipCompileArgs {
     #[clap(required(true))]
     src_file: Vec<PathBuf>,
 
-    /// Constrain versions using the given constraints files.
+    /// Constrain versions using the given requirements files.
+    ///
+    /// Constraints files are `requirements.txt`-like files that only control the _version_ of a
+    /// requirement that's installed. However, including a package in a constraints file will _not_
+    /// trigger the installation of that package.
+    ///
+    /// This is equivalent to pip's `--constraint` option.
     #[clap(short, long)]
     constraint: Vec<PathBuf>,
 
+    /// Override versions using the given requirements files.
+    ///
+    /// Overrides files are `requirements.txt`-like files that force a specific version of a
+    /// requirement to be installed, regardless of the requirements declared by any constituent
+    /// package, and regardless of whether this would be considered an invalid resolution.
+    ///
+    /// While constraints are _additive_, in that they're combined with the requirements of the
+    /// constituent packages, overrides are _absolute_, in that they completely replace the
+    /// requirements of the constituent packages.
+    #[clap(long)]
+    r#override: Vec<PathBuf>,
+
     /// Include optional dependencies in the given extra group name; may be provided more than once.
     #[clap(long, conflicts_with = "all_extras", value_parser = extra_name_with_clap_error)]
     extra: Vec<ExtraName>,
@@ -228,10 +246,28 @@ struct PipInstallArgs {
     #[clap(short, long, group = "sources")]
     requirement: Vec<PathBuf>,
 
-    /// Constrain versions using the given constraints files.
+    /// Constrain versions using the given requirements files.
+    ///
+    /// Constraints files are `requirements.txt`-like files that only control the _version_ of a
+    /// requirement that's installed. However, including a package in a constraints file will _not_
+    /// trigger the installation of that package.
+    ///
+    /// This is equivalent to pip's `--constraint` option.
     #[clap(short, long)]
     constraint: Vec<PathBuf>,
 
+    /// Override versions using the given requirements files.
+    ///
+    /// Overrides files are `requirements.txt`-like files that force a specific version of a
+    /// requirement to be installed, regardless of the requirements declared by any constituent
+    /// package, and regardless of whether this would be considered an invalid resolution.
+    ///
+    /// While constraints are _additive_, in that they're combined with the requirements of the
+    /// constituent packages, overrides are _absolute_, in that they completely replace the
+    /// requirements of the constituent packages.
+    #[clap(long)]
+    r#override: Vec<PathBuf>,
+
     /// Include optional dependencies in the given extra group name; may be provided more than once.
     #[clap(long, conflicts_with = "all_extras", value_parser = extra_name_with_clap_error)]
     extra: Vec<ExtraName>,
@@ -383,6 +419,11 @@ async fn inner() -> Result<ExitStatus> {
                 .into_iter()
                 .map(RequirementsSource::from)
                 .collect::<Vec<_>>();
+            let overrides = args
+                .r#override
+                .into_iter()
+                .map(RequirementsSource::from)
+                .collect::<Vec<_>>();
             let index_urls =
                 IndexUrls::from_args(args.index_url, args.extra_index_url, args.no_index);
             let extras = if args.all_extras {
@@ -395,6 +436,7 @@ async fn inner() -> Result<ExitStatus> {
             commands::pip_compile(
                 &requirements,
                 &constraints,
+                &overrides,
                 extras,
                 args.output_file.as_deref(),
                 args.resolution.unwrap_or_default(),
@@ -441,6 +483,11 @@ async fn inner() -> Result<ExitStatus> {
                 .into_iter()
                 .map(RequirementsSource::from)
                 .collect::<Vec<_>>();
+            let overrides = args
+                .r#override
+                .into_iter()
+                .map(RequirementsSource::from)
+                .collect::<Vec<_>>();
             let index_urls =
                 IndexUrls::from_args(args.index_url, args.extra_index_url, args.no_index);
             let extras = if args.all_extras {
@@ -454,6 +501,7 @@ async fn inner() -> Result<ExitStatus> {
             commands::pip_install(
                 &requirements,
                 &constraints,
+                &overrides,
                 &extras,
                 args.resolution.unwrap_or_default(),
                 args.prerelease.unwrap_or_default(),

crates/puffin-cli/src/requirements.rs

@@ -64,6 +64,8 @@ pub(crate) struct RequirementsSpecification {
     pub(crate) requirements: Vec<Requirement>,
     /// The constraints for the project.
     pub(crate) constraints: Vec<Requirement>,
+    /// The overrides for the project.
+    pub(crate) overrides: Vec<Requirement>,
     /// The extras used to collect requirements.
     pub(crate) extras: FxHashSet<ExtraName>,
 }
@@ -82,6 +84,7 @@ impl RequirementsSpecification {
                     project: None,
                     requirements: vec![requirement],
                     constraints: vec![],
+                    overrides: vec![],
                     extras: FxHashSet::default(),
                 }
             }
@@ -95,6 +98,7 @@ impl RequirementsSpecification {
                         .map(|entry| entry.requirement)
                         .collect(),
                     constraints: requirements_txt.constraints.into_iter().collect(),
+                    overrides: vec![],
                     extras: FxHashSet::default(),
                 }
             }
@@ -131,6 +135,7 @@ impl RequirementsSpecification {
                     project: project_name,
                     requirements,
                     constraints: vec![],
+                    overrides: vec![],
                     extras: used_extras,
                 }
             }
@@ -141,6 +146,7 @@ impl RequirementsSpecification {
     pub(crate) fn from_sources(
         requirements: &[RequirementsSource],
         constraints: &[RequirementsSource],
+        overrides: &[RequirementsSource],
         extras: &ExtrasSpecification,
     ) -> Result<Self> {
         let mut spec = Self::default();
@@ -152,6 +158,7 @@ impl RequirementsSpecification {
             let source = Self::from_source(source, extras)?;
             spec.requirements.extend(source.requirements);
             spec.constraints.extend(source.constraints);
+            spec.overrides.extend(source.overrides);
             spec.extras.extend(source.extras);
 
             // Use the first project name discovered
@@ -160,18 +167,27 @@ impl RequirementsSpecification {
             }
         }
 
-        // Read all constraints, treating both requirements _and_ constraints as constraints.
+        // Read all constraints, treating _everything_ as a constraint.
         for source in constraints {
             let source = Self::from_source(source, extras)?;
             spec.constraints.extend(source.requirements);
             spec.constraints.extend(source.constraints);
+            spec.constraints.extend(source.overrides);
+        }
+
+        // Read all overrides, treating both requirements _and_ constraints as overrides.
+        for source in overrides {
+            let source = Self::from_source(source, extras)?;
+            spec.overrides.extend(source.requirements);
+            spec.overrides.extend(source.constraints);
+            spec.overrides.extend(source.overrides);
         }
 
         Ok(spec)
     }
 
     /// Read the requirements from a set of sources.
     pub(crate) fn requirements(requirements: &[RequirementsSource]) -> Result<Vec<Requirement>> {
-        Ok(Self::from_sources(requirements, &[], &ExtrasSpecification::None)?.requirements)
+        Ok(Self::from_sources(requirements, &[], &[], &ExtrasSpecification::None)?.requirements)
     }
 }

crates/puffin-cli/tests/pip_compile.rs

@@ -1606,3 +1606,117 @@ fn compile_yanked_version_indirect() -> Result<()> {
 
     Ok(())
 }
+
+/// Flask==3.0.0 depends on Werkzeug>=3.0.0. Demonstrate that we can override this
+/// requirement with an incompatible version.
+#[test]
+fn override_dependency() -> Result<()> {
+    let temp_dir = TempDir::new()?;
+    let cache_dir = TempDir::new()?;
+    let venv = create_venv_py312(&temp_dir, &cache_dir);
+
+    let requirements_in = temp_dir.child("requirements.in");
+    requirements_in.write_str("flask==3.0.0")?;
+
+    let overrides_txt = temp_dir.child("overrides.txt");
+    overrides_txt.write_str("werkzeug==2.3.0")?;
+
+    insta::with_settings!({
+        filters => INSTA_FILTERS.to_vec()
+    }, {
+        assert_cmd_snapshot!(Command::new(get_cargo_bin(BIN_NAME))
+            .arg("pip-compile")
+            .arg("requirements.in")
+            .arg("--override")
+            .arg("overrides.txt")
+            .arg("--cache-dir")
+            .arg(cache_dir.path())
+            .arg("--exclude-newer")
+            .arg(EXCLUDE_NEWER)
+            .env("VIRTUAL_ENV", venv.as_os_str())
+            .current_dir(&temp_dir), @r###"
+        success: true
+        exit_code: 0
+        ----- stdout -----
+        # This file was autogenerated by Puffin v0.0.1 via the following command:
+        #    puffin pip-compile requirements.in --override overrides.txt --cache-dir [CACHE_DIR]
+        blinker==1.7.0
+            # via flask
+        click==8.1.7
+            # via flask
+        flask==3.0.0
+        itsdangerous==2.1.2
+            # via flask
+        jinja2==3.1.2
+            # via flask
+        markupsafe==2.1.3
+            # via
+            #   jinja2
+            #   werkzeug
+        werkzeug==2.3.0
+            # via flask
+
+        ----- stderr -----
+        Resolved 7 packages in [TIME]
+        "###);
+    });
+
+    Ok(())
+}
+
+/// Black==23.10.1 depends on tomli>=1.1.0 for Python versions below 3.11. Demonstrate that we can
+/// override it with a multi-line override.
+#[test]
+fn override_multi_dependency() -> Result<()> {
+    let temp_dir = TempDir::new()?;
+    let cache_dir = TempDir::new()?;
+    let venv = create_venv_py312(&temp_dir, &cache_dir);
+
+    let requirements_in = temp_dir.child("requirements.in");
+    requirements_in.write_str("black==23.10.1")?;
+
+    let overrides_txt = temp_dir.child("overrides.txt");
+    overrides_txt.write_str(
+        "tomli>=1.1.0; python_version >= '3.11'\ntomli<1.0.0; python_version < '3.11'",
+    )?;
+
+    insta::with_settings!({
+        filters => INSTA_FILTERS.to_vec()
+    }, {
+        assert_cmd_snapshot!(Command::new(get_cargo_bin(BIN_NAME))
+            .arg("pip-compile")
+            .arg("requirements.in")
+            .arg("--override")
+            .arg("overrides.txt")
+            .arg("--cache-dir")
+            .arg(cache_dir.path())
+            .arg("--exclude-newer")
+            .arg(EXCLUDE_NEWER)
+            .env("VIRTUAL_ENV", venv.as_os_str())
+            .current_dir(&temp_dir), @r###"
+        success: true
+        exit_code: 0
+        ----- stdout -----
+        # This file was autogenerated by Puffin v0.0.1 via the following command:
+        #    puffin pip-compile requirements.in --override overrides.txt --cache-dir [CACHE_DIR]
+        black==23.10.1
+        click==8.1.7
+            # via black
+        mypy-extensions==1.0.0
+            # via black
+        packaging==23.2
+            # via black
+        pathspec==0.11.2
+            # via black
+        platformdirs==4.0.0
+            # via black
+        tomli==2.0.1
+            # via black
+
+        ----- stderr -----
+        Resolved 7 packages in [TIME]
+        "###);
+    });
+
+    Ok(())
+}

crates/puffin-dev/src/resolve_cli.rs

@@ -60,12 +60,7 @@ pub(crate) async fn resolve_cli(args: ResolveCliArgs) -> Result<()> {
     // Copied from `BuildDispatch`
     let tags = Tags::from_interpreter(venv.interpreter())?;
     let resolver = Resolver::new(
-        Manifest::new(
-            args.requirements.clone(),
-            Vec::default(),
-            Vec::default(),
-            None,
-        ),
+        Manifest::simple(args.requirements.clone()),
         ResolutionOptions::default(),
         venv.interpreter().markers(),
         &tags,

crates/puffin-dispatch/src/lib.rs

@@ -92,7 +92,7 @@ impl BuildContext for BuildDispatch {
         Box::pin(async {
             let tags = Tags::from_interpreter(&self.interpreter)?;
             let resolver = Resolver::new(
-                Manifest::new(requirements.to_vec(), Vec::default(), Vec::default(), None),
+                Manifest::simple(requirements.to_vec()),
                 self.options,
                 self.interpreter.markers(),
                 &tags,