Support multiple extras in universal pip compile output

Author: charliermarsh

crates/uv-resolver/src/resolution/display.rs

@@ -5,7 +5,9 @@ use petgraph::visit::EdgeRef;
 use petgraph::Direction;
 use rustc_hash::{FxBuildHasher, FxHashMap};
 
-use uv_distribution_types::{DistributionMetadata, Name, SourceAnnotation, SourceAnnotations};
+use uv_distribution_types::{
+    DistributionMetadata, Name, SourceAnnotation, SourceAnnotations, VersionId,
+};
 use uv_normalize::PackageName;
 use uv_pep508::MarkerTree;
 
@@ -44,15 +46,6 @@ enum DisplayResolutionGraphNode<'dist> {
     Dist(RequirementsTxtDist<'dist>),
 }
 
-impl DisplayResolutionGraphNode<'_> {
-    fn markers(&self) -> &MarkerTree {
-        match self {
-            DisplayResolutionGraphNode::Root => &MarkerTree::TRUE,
-            DisplayResolutionGraphNode::Dist(dist) => dist.markers,
-        }
-    }
-}
-
 impl<'a> DisplayResolutionGraph<'a> {
     /// Create a new [`DisplayResolutionGraph`] for the given graph.
     #[allow(clippy::fn_params_excessive_bools)]
@@ -156,9 +149,12 @@ impl std::fmt::Display for DisplayResolutionGraph<'_> {
             |_index, _edge| (),
         );
 
-        // Reduce the graph, such that all nodes for a single package are combined, regardless of
-        // the extras.
-        let petgraph = combine_extras(&petgraph);
+        // Reduce the graph, removing or combining extras for a given package.
+        let petgraph = if self.include_extras {
+            combine_extras(&petgraph)
+        } else {
+            strip_extras(&petgraph)
+        };
 
         // Collect all packages.
         let mut nodes = petgraph
@@ -181,11 +177,7 @@ impl std::fmt::Display for DisplayResolutionGraph<'_> {
         for (index, node) in nodes {
             // Display the node itself.
             let mut line = node
-                .to_requirements_txt(
-                    &self.resolution.requires_python,
-                    self.include_extras,
-                    self.include_markers,
-                )
+                .to_requirements_txt(&self.resolution.requires_python, self.include_markers)
                 .to_string();
 
             // Display the distribution hashes, if any.
@@ -320,13 +312,22 @@ type RequirementsTxtGraph<'dist> =
     petgraph::graph::Graph<RequirementsTxtDist<'dist>, (), petgraph::Directed>;
 
 /// Reduce the graph, such that all nodes for a single package are combined, regardless of
-/// the extras.
+/// the extras, as long as they have the same version and markers.
 ///
 /// For example, `flask` and `flask[dotenv]` should be reduced into a single `flask[dotenv]`
 /// node.
 ///
+/// If the extras have different markers, they'll be treated as separate nodes. For example,
+/// `flask[dotenv] ; sys_platform == "win32"` and `flask[async] ; sys_platform == "linux"`
+/// would _not_ be combined.
+///
 /// We also remove the root node, to simplify the graph structure.
 fn combine_extras<'dist>(graph: &IntermediatePetGraph<'dist>) -> RequirementsTxtGraph<'dist> {
+    /// Return the key for a node.
+    fn version_marker(dist: &RequirementsTxtDist) -> (VersionId, MarkerTree) {
+        (dist.version_id(), dist.markers.clone())
+    }
+
     let mut next = RequirementsTxtGraph::with_capacity(graph.node_count(), graph.edge_count());
     let mut inverse = FxHashMap::with_capacity_and_hasher(graph.node_count(), FxBuildHasher);
 
@@ -338,40 +339,68 @@ fn combine_extras<'dist>(graph: &IntermediatePetGraph<'dist>) -> RequirementsTxt
 
         // In the `requirements.txt` output, we want a flat installation list, so we need to use
         // the reachability markers instead of the edge markers.
-        // We use the markers of the base package: We know that each virtual extra package has an
-        // edge to the base package, so we know that base package markers are more general than the
-        // extra package markers (the extra package markers are a subset of the base package
-        // markers).
-        if let Some(index) = inverse.get(&dist.version_id()) {
-            let node: &mut RequirementsTxtDist = &mut next[*index];
-            node.extras.extend(dist.extras.iter().cloned());
-            node.extras.sort_unstable();
-            node.extras.dedup();
-        } else {
-            let version_id = dist.version_id();
-            let dist = dist.clone();
-            let index = next.add_node(dist);
-            inverse.insert(version_id, index);
+        match inverse.entry(version_marker(dist)) {
+            std::collections::hash_map::Entry::Occupied(entry) => {
+                let index = *entry.get();
+                let node: &mut RequirementsTxtDist = &mut next[index];
+                node.extras.extend(dist.extras.iter().cloned());
+                node.extras.sort_unstable();
+                node.extras.dedup();
+            }
+            std::collections::hash_map::Entry::Vacant(entry) => {
+                let index = next.add_node(dist.clone());
+                entry.insert(index);
+            }
         }
     }
 
-    // Verify that the package markers are more general than the extra markers.
-    if cfg!(debug_assertions) {
-        for index in graph.node_indices() {
-            let DisplayResolutionGraphNode::Dist(dist) = &graph[index] else {
-                continue;
-            };
-            let combined_markers = next[inverse[&dist.version_id()]].markers.clone();
-            let mut package_markers = combined_markers.clone();
-            package_markers.or(graph[index].markers().clone());
-            assert_eq!(
-                package_markers,
-                combined_markers,
-                "{} {:?} {:?}",
-                dist.version_id(),
-                dist.extras,
-                dist.markers.try_to_string()
-            );
+    // Re-add the edges to the reduced graph.
+    for edge in graph.edge_indices() {
+        let (source, target) = graph.edge_endpoints(edge).unwrap();
+        let DisplayResolutionGraphNode::Dist(source_node) = &graph[source] else {
+            continue;
+        };
+        let DisplayResolutionGraphNode::Dist(target_node) = &graph[target] else {
+            continue;
+        };
+        let source = inverse[&version_marker(source_node)];
+        let target = inverse[&version_marker(target_node)];
+
+        next.update_edge(source, target, ());
+    }
+
+    next
+}
+
+/// Reduce the graph, such that all nodes for a single package are combined, with extras
+/// removed.
+///
+/// For example, `flask`, `flask[async]`, and `flask[dotenv]` should be reduced into a single
+/// `flask` node, with a conjunction of their markers.
+///
+/// We also remove the root node, to simplify the graph structure.
+fn strip_extras<'dist>(graph: &IntermediatePetGraph<'dist>) -> RequirementsTxtGraph<'dist> {
+    let mut next = RequirementsTxtGraph::with_capacity(graph.node_count(), graph.edge_count());
+    let mut inverse = FxHashMap::with_capacity_and_hasher(graph.node_count(), FxBuildHasher);
+
+    // Re-add the nodes to the reduced graph.
+    for index in graph.node_indices() {
+        let DisplayResolutionGraphNode::Dist(dist) = &graph[index] else {
+            continue;
+        };
+
+        // In the `requirements.txt` output, we want a flat installation list, so we need to use
+        // the reachability markers instead of the edge markers.
+        match inverse.entry(dist.version_id()) {
+            std::collections::hash_map::Entry::Occupied(entry) => {
+                let index = *entry.get();
+                let node: &mut RequirementsTxtDist = &mut next[index];
+                node.extras.clear();
+            }
+            std::collections::hash_map::Entry::Vacant(entry) => {
+                let index = next.add_node(dist.clone());
+                entry.insert(index);
+            }
         }
     }
 

crates/uv-resolver/src/resolution/requirements_txt.rs

@@ -35,7 +35,6 @@ impl<'dist> RequirementsTxtDist<'dist> {
     pub(crate) fn to_requirements_txt(
         &self,
         requires_python: &RequiresPython,
-        include_extras: bool,
         include_markers: bool,
     ) -> Cow<str> {
         // If the URL is editable, write it as an editable requirement.
@@ -106,7 +105,7 @@ impl<'dist> RequirementsTxtDist<'dist> {
             }
         }
 
-        if self.extras.is_empty() || !include_extras {
+        if self.extras.is_empty() {
             if let Some(markers) = SimplifiedMarkerTree::new(requires_python, self.markers.clone())
                 .try_to_string()
                 .filter(|_| include_markers)
@@ -141,6 +140,8 @@ impl<'dist> RequirementsTxtDist<'dist> {
         }
     }
 
+    /// Convert the [`RequirementsTxtDist`] to a comparator that can be used to sort the requirements
+    /// in a `requirements.txt` file.
     pub(crate) fn to_comparator(&self) -> RequirementsTxtComparator {
         if self.dist.is_editable() {
             if let VersionOrUrlRef::Url(url) = self.dist.version_or_url() {
@@ -153,12 +154,14 @@ impl<'dist> RequirementsTxtDist<'dist> {
                 name: self.name(),
                 version: self.version,
                 url: Some(url.verbatim()),
+                extras: &self.extras,
             }
         } else {
             RequirementsTxtComparator::Name {
                 name: self.name(),
                 version: self.version,
                 url: None,
+                extras: &self.extras,
             }
         }
     }
@@ -178,15 +181,18 @@ impl<'dist> RequirementsTxtDist<'dist> {
     }
 }
 
+/// A comparator for sorting requirements in a `requirements.txt` file.
 #[derive(Debug, Clone, PartialEq, Eq, PartialOrd, Ord)]
 pub(crate) enum RequirementsTxtComparator<'a> {
+    /// Sort by URL for editable requirements.
     Url(Cow<'a, str>),
     /// In universal mode, we can have multiple versions for a package, so we track the version and
     /// the URL (for non-index packages) to have a stable sort for those, too.
     Name {
         name: &'a PackageName,
         version: &'a Version,
         url: Option<Cow<'a, str>>,
+        extras: &'a [ExtraName],
     },
 }