Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Support multiple extras in universal pip compile output #8960

Merged
merged 1 commit into from
Nov 8, 2024
Merged

Support multiple extras in universal pip compile output #8960

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
129 changes: 79 additions & 50 deletions crates/uv-resolver/src/resolution/display.rs
View file
Open in desktop
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);
}
}
}
Copy link
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I could make this a trait to DRY up the above but... I don't know, it seems like a lot of complexity.


Expand Down
10 changes: 8 additions & 2 deletions crates/uv-resolver/src/resolution/requirements_txt.rs
View file
Open in desktop
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
Loading