Apply extra to overrides and constraints (#4829)

This is an attempt to solve https://github.com/astral-sh/uv/issues/ by
applying the extra marker of the requirement to overrides and
constraints.

Say in `a` we have a requirements
```
b==1; python_version < "3.10"
c==1; extra == "feature"
```

and overrides
```
b==2; python_version < "3.10"
b==3; python_version >= "3.10"
c==2; python_version < "3.10"
c==3; python_version >= "3.10"
```

Our current strategy is to discard the markers in the original
requirements. This means that on 3.12 for `a` we install `b==3`, but it
also means that we add `c` to `a` without `a[feature]`, causing #4826.
With this PR, the new requirement become,

```
b==2; python_version < "3.10"
b==3; python_version >= "3.10"
c==2; python_version < "3.10" and extra == "feature"
c==3; python_version >= "3.10" and extra == "feature"
```

allowing to override markers while preserving optional dependencies as
such.

Fixes #4826

crates/pep508-rs/src/marker.rs

@@ -1897,6 +1897,28 @@ impl MarkerTree {
             }
         }
     }
+
+    /// Find a top level `extra == "..."` expression.
+    ///
+    /// ASSUMPTION: There is one `extra = "..."`, and it's either the only marker or part of the
+    /// main conjunction.
+    pub fn top_level_extra(&self) -> Option<&MarkerExpression> {
+        match &self {
+            MarkerTree::Expression(extra_expression @ MarkerExpression::Extra { .. }) => {
+                Some(extra_expression)
+            }
+            MarkerTree::And(and) => and.iter().find_map(|marker| {
+                if let MarkerTree::Expression(extra_expression @ MarkerExpression::Extra { .. }) =
+                    marker
+                {
+                    Some(extra_expression)
+                } else {
+                    None
+                }
+            }),
+            MarkerTree::Expression(_) | MarkerTree::Or(_) => None,
+        }
+    }
 }
 
 impl Display for MarkerTree {

crates/uv-configuration/src/constraints.rs

@@ -1,6 +1,8 @@
-use rustc_hash::{FxBuildHasher, FxHashMap};
-
+use either::Either;
+use pep508_rs::MarkerTree;
 use pypi_types::Requirement;
+use rustc_hash::{FxBuildHasher, FxHashMap};
+use std::borrow::Cow;
 use uv_normalize::PackageName;
 
 /// A set of constraints for a set of requirements.
@@ -36,16 +38,49 @@ impl Constraints {
     }
 
     /// Apply the constraints to a set of requirements.
+    ///
+    /// NB: Change this method together with [`Overrides::apply`].
     pub fn apply<'a>(
         &'a self,
-        requirements: impl IntoIterator<Item = &'a Requirement>,
-    ) -> impl Iterator<Item = &Requirement> {
+        requirements: impl IntoIterator<Item = Cow<'a, Requirement>>,
+    ) -> impl Iterator<Item = Cow<'a, Requirement>> {
         requirements.into_iter().flat_map(|requirement| {
-            std::iter::once(requirement).chain(
-                self.get(&requirement.name)
-                    .into_iter()
-                    .flat_map(|constraints| constraints.iter()),
-            )
+            let Some(constraints) = self.get(&requirement.name) else {
+                // Case 1: No constraint(s).
+                return Either::Left(std::iter::once(requirement));
+            };
+
+            // ASSUMPTION: There is one `extra = "..."`, and it's either the only marker or part
+            // of the main conjunction.
+            let Some(extra_expression) = requirement
+                .marker
+                .as_ref()
+                .and_then(|marker| marker.top_level_extra())
+                .cloned()
+            else {
+                // Case 2: A non-optional dependency with constraint(s).
+                return Either::Right(Either::Right(
+                    std::iter::once(requirement).chain(constraints.iter().map(Cow::Borrowed)),
+                ));
+            };
+
+            // Case 3: An optional dependency with constraint(s).
+            //
+            // When the original requirement is an optional dependency, the constraint(s) need to
+            // be optional for the same extra, otherwise we activate extras that should be inactive.
+            Either::Right(Either::Left(std::iter::once(requirement).chain(
+                constraints.iter().cloned().map(move |constraint| {
+                    // Add the extra to the override marker.
+                    let mut joint_marker = MarkerTree::Expression(extra_expression.clone());
+                    if let Some(marker) = &constraint.marker {
+                        joint_marker.and(marker.clone());
+                    }
+                    Cow::Owned(Requirement {
+                        marker: Some(joint_marker.clone()),
+                        ..constraint
+                    })
+                }),
+            )))
         })
     }
 }

crates/uv-configuration/src/overrides.rs

@@ -1,6 +1,9 @@
+use std::borrow::Cow;
+
 use either::Either;
 use rustc_hash::{FxBuildHasher, FxHashMap};
 
+use pep508_rs::MarkerTree;
 use pypi_types::Requirement;
 use uv_normalize::PackageName;
 
@@ -33,16 +36,44 @@ impl Overrides {
     }
 
     /// Apply the overrides to a set of requirements.
+    ///
+    /// NB: Change this method together with [`Constraints::apply`].
     pub fn apply<'a>(
         &'a self,
         requirements: impl IntoIterator<Item = &'a Requirement>,
-    ) -> impl Iterator<Item = &Requirement> {
+    ) -> impl Iterator<Item = Cow<'a, Requirement>> {
         requirements.into_iter().flat_map(|requirement| {
-            if let Some(overrides) = self.get(&requirement.name) {
-                Either::Left(overrides.iter())
-            } else {
-                Either::Right(std::iter::once(requirement))
-            }
+            let Some(overrides) = self.get(&requirement.name) else {
+                // Case 1: No override(s).
+                return Either::Left(std::iter::once(Cow::Borrowed(requirement)));
+            };
+
+            // ASSUMPTION: There is one `extra = "..."`, and it's either the only marker or part
+            // of the main conjunction.
+            let Some(extra_expression) = requirement
+                .marker
+                .as_ref()
+                .and_then(|marker| marker.top_level_extra())
+            else {
+                // Case 2: A non-optional dependency with override(s).
+                return Either::Right(Either::Right(overrides.iter().map(Cow::Borrowed)));
+            };
+
+            // Case 3: An optional dependency with override(s).
+            //
+            // When the original requirement is an optional dependency, the override(s) need to
+            // be optional for the same extra, otherwise we activate extras that should be inactive.
+            Either::Right(Either::Left(overrides.iter().map(|override_requirement| {
+                // Add the extra to the override marker.
+                let mut joint_marker = MarkerTree::Expression(extra_expression.clone());
+                if let Some(marker) = &override_requirement.marker {
+                    joint_marker.and(marker.clone());
+                }
+                Cow::Owned(Requirement {
+                    marker: Some(joint_marker.clone()),
+                    ..override_requirement.clone()
+                })
+            })))
         })
     }
 }

crates/uv-requirements/src/lookahead.rs

@@ -109,7 +109,7 @@ impl<'a, Context: BuildContext> LookaheadResolver<'a, Context> {
             .constraints
             .apply(self.overrides.apply(self.requirements))
             .filter(|requirement| requirement.evaluate_markers(markers, &[]))
-            .cloned()
+            .map(|requirement| (*requirement).clone())
             .collect();
 
         while !queue.is_empty() || !futures.is_empty() {
@@ -128,7 +128,7 @@ impl<'a, Context: BuildContext> LookaheadResolver<'a, Context> {
                         .apply(self.overrides.apply(lookahead.requirements()))
                     {
                         if requirement.evaluate_markers(markers, lookahead.extras()) {
-                            queue.push_back(requirement.clone());
+                            queue.push_back((*requirement).clone());
                         }
                     }
                     results.push(lookahead);

crates/uv-resolver/src/manifest.rs

@@ -1,4 +1,5 @@
 use either::Either;
+use std::borrow::Cow;
 
 use pep508_rs::MarkerEnvironment;
 use pypi_types::Requirement;
@@ -97,7 +98,7 @@ impl Manifest {
         &'a self,
         markers: Option<&'a MarkerEnvironment>,
         mode: DependencyMode,
-    ) -> impl Iterator<Item = &Requirement> + 'a {
+    ) -> impl Iterator<Item = Cow<'a, Requirement>> + 'a {
         self.requirements_no_overrides(markers, mode)
             .chain(self.overrides(markers, mode))
     }
@@ -107,7 +108,7 @@ impl Manifest {
         &'a self,
         markers: Option<&'a MarkerEnvironment>,
         mode: DependencyMode,
-    ) -> impl Iterator<Item = &Requirement> + 'a {
+    ) -> impl Iterator<Item = Cow<'a, Requirement>> + 'a {
         match mode {
             // Include all direct and transitive requirements, with constraints and overrides applied.
             DependencyMode::Transitive => Either::Left(
@@ -128,14 +129,15 @@ impl Manifest {
                     .chain(
                         self.constraints
                             .requirements()
-                            .filter(move |requirement| requirement.evaluate_markers(markers, &[])),
+                            .filter(move |requirement| requirement.evaluate_markers(markers, &[]))
+                            .map(Cow::Borrowed),
                     ),
             ),
             // Include direct requirements, with constraints and overrides applied.
             DependencyMode::Direct => Either::Right(
                 self.overrides
                     .apply(&self.requirements)
-                    .chain(self.constraints.requirements())
+                    .chain(self.constraints.requirements().map(Cow::Borrowed))
                     .filter(move |requirement| requirement.evaluate_markers(markers, &[])),
             ),
         }
@@ -146,19 +148,21 @@ impl Manifest {
         &'a self,
         markers: Option<&'a MarkerEnvironment>,
         mode: DependencyMode,
-    ) -> impl Iterator<Item = &Requirement> + 'a {
+    ) -> impl Iterator<Item = Cow<'a, Requirement>> + 'a {
         match mode {
             // Include all direct and transitive requirements, with constraints and overrides applied.
             DependencyMode::Transitive => Either::Left(
                 self.overrides
                     .requirements()
-                    .filter(move |requirement| requirement.evaluate_markers(markers, &[])),
+                    .filter(move |requirement| requirement.evaluate_markers(markers, &[]))
+                    .map(Cow::Borrowed),
             ),
             // Include direct requirements, with constraints and overrides applied.
             DependencyMode::Direct => Either::Right(
                 self.overrides
                     .requirements()
-                    .filter(move |requirement| requirement.evaluate_markers(markers, &[])),
+                    .filter(move |requirement| requirement.evaluate_markers(markers, &[]))
+                    .map(Cow::Borrowed),
             ),
         }
     }
@@ -177,7 +181,7 @@ impl Manifest {
         &'a self,
         markers: Option<&'a MarkerEnvironment>,
         mode: DependencyMode,
-    ) -> impl Iterator<Item = &PackageName> + 'a {
+    ) -> impl Iterator<Item = Cow<'a, PackageName>> + 'a {
         match mode {
             // Include direct requirements, dependencies of editables, and transitive dependencies
             // of local packages.
@@ -197,7 +201,10 @@ impl Manifest {
                             .apply(&self.requirements)
                             .filter(move |requirement| requirement.evaluate_markers(markers, &[])),
                     )
-                    .map(|requirement| &requirement.name),
+                    .map(|requirement| match requirement {
+                        Cow::Borrowed(requirement) => Cow::Borrowed(&requirement.name),
+                        Cow::Owned(requirement) => Cow::Owned(requirement.name),
+                    }),
             ),
 
             // Restrict to the direct requirements.
@@ -205,7 +212,10 @@ impl Manifest {
                 self.overrides
                     .apply(self.requirements.iter())
                     .filter(move |requirement| requirement.evaluate_markers(markers, &[]))
-                    .map(|requirement| &requirement.name),
+                    .map(|requirement| match requirement {
+                        Cow::Borrowed(requirement) => Cow::Borrowed(&requirement.name),
+                        Cow::Owned(requirement) => Cow::Owned(requirement.name),
+                    }),
             ),
         }
     }
@@ -217,7 +227,7 @@ impl Manifest {
     pub fn apply<'a>(
         &'a self,
         requirements: impl IntoIterator<Item = &'a Requirement>,
-    ) -> impl Iterator<Item = &Requirement> {
+    ) -> impl Iterator<Item = Cow<'a, Requirement>> {
         self.constraints.apply(self.overrides.apply(requirements))
     }
 

crates/uv-resolver/src/resolution_mode.rs

@@ -46,7 +46,7 @@ impl ResolutionStrategy {
             ResolutionMode::LowestDirect => Self::LowestDirect(
                 manifest
                     .user_requirements(markers, dependencies)
-                    .cloned()
+                    .map(|requirement| (*requirement).clone())
                     .collect(),
             ),
         }

crates/uv-resolver/src/resolver/mod.rs

@@ -1477,7 +1477,7 @@ impl<InstalledPackages: InstalledPackagesProvider> ResolverState<InstalledPackag
                 true
             })
             .flat_map(move |requirement| {
-                iter::once(Cow::Borrowed(requirement)).chain(
+                iter::once(requirement.clone()).chain(
                     self.constraints
                         .get(&requirement.name)
                         .into_iter()

crates/uv/tests/pip_compile.rs

@@ -3154,6 +3154,37 @@ fn override_multi_dependency() -> Result<()> {
     Ok(())
 }
 
+/// `urllib3==2.2.2` has an optional dependency on `pysocks!=1.5.7,<2.0,>=1.5.6; extra == 'socks'`,
+/// So we shouldn't apply the `pysocks==1.7.1` override without the `socks` extra.
+#[test]
+fn dont_add_override_for_non_activated_extra() -> Result<()> {
+    let context = TestContext::new("3.12");
+    let requirements_in = context.temp_dir.child("requirements.in");
+    requirements_in.write_str("urllib3==2.2.1")?;
+
+    let overrides_txt = context.temp_dir.child("overrides.txt");
+    overrides_txt.write_str("pysocks==1.7.1")?;
+
+    uv_snapshot!(context.pip_compile()
+        .arg("requirements.in")
+        .arg("--override")
+        .arg("overrides.txt"), @r###"
+    success: true
+    exit_code: 0
+    ----- stdout -----
+    # This file was autogenerated by uv via the following command:
+    #    uv pip compile --cache-dir [CACHE_DIR] requirements.in --override overrides.txt
+    urllib3==2.2.1
+        # via -r requirements.in
+
+    ----- stderr -----
+    Resolved 1 package in [TIME]
+    "###
+    );
+
+    Ok(())
+}
+
 /// Check how invalid `tool.uv.override-dependencies` is handled in `pyproject.toml`.
 #[test]
 fn override_dependency_from_workspace_invalid_syntax() -> Result<()> {