Skip to content

Commit

Permalink
Support multiple extras in universal pip compile output (#8960)
Browse files Browse the repository at this point in the history
## Summary

We were making some incorrect assumptions in the extra-merging code for
universal `pip compile`. This PR corrects those assumptions and adds a
bunch of additional tests.

Closes #8915.
  • Loading branch information
charliermarsh authored and konstin committed Nov 10, 2024
1 parent 55a7ebf commit 0b15684
Show file tree
Hide file tree
Showing 3 changed files with 479 additions and 52 deletions.
129 changes: 79 additions & 50 deletions crates/uv-resolver/src/resolution/display.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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;

Expand Down Expand Up @@ -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)]
Expand Down Expand Up @@ -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
Expand All @@ -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.
Expand Down Expand Up @@ -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);

Expand All @@ -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);
}
}
}

Expand Down
10 changes: 8 additions & 2 deletions crates/uv-resolver/src/resolution/requirements_txt.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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.
Expand Down Expand Up @@ -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)
Expand Down Expand Up @@ -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() {
Expand All @@ -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,
}
}
}
Expand All @@ -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],
},
}

Expand Down
Loading

0 comments on commit 0b15684

Please sign in to comment.