From 19c91e7dacda9f78e4d8cc576177f52fb1516fdf Mon Sep 17 00:00:00 2001 From: Charlie Marsh Date: Wed, 29 May 2024 11:42:49 -0400 Subject: [PATCH] Create distinct graph nodes for each package extra (#3908) ## Summary Today, we represent each package as a single node in the graph, and combine all the extras. This is helpful for the `requirements.txt`-style resolution, in which we want to show each a single line for each package with the extras combined into a single array. This PR modifies the representation to instead use a separate node for each (package, extra) pair. We then reduce into the previous format when printing in the `requirements.txt`-style format, so there shouldn't be any user-facing changes here. --- crates/uv-resolver/src/resolution/display.rs | 57 ++++++- crates/uv-resolver/src/resolution/graph.rs | 150 ++++++------------- 2 files changed, 97 insertions(+), 110 deletions(-) diff --git a/crates/uv-resolver/src/resolution/display.rs b/crates/uv-resolver/src/resolution/display.rs index dbda99f50..27dd727ed 100644 --- a/crates/uv-resolver/src/resolution/display.rs +++ b/crates/uv-resolver/src/resolution/display.rs @@ -1,12 +1,15 @@ use std::collections::BTreeSet; +use std::hash::BuildHasherDefault; use owo_colors::OwoColorize; use petgraph::visit::EdgeRef; use petgraph::Direction; +use rustc_hash::FxHashMap; use distribution_types::{Name, SourceAnnotations}; use uv_normalize::PackageName; +use crate::resolution::AnnotatedDist; use crate::ResolutionGraph; /// A [`std::fmt::Display`] implementation for the resolution graph. @@ -77,13 +80,53 @@ impl<'a> DisplayResolutionGraph<'a> { /// Write the graph in the `{name}=={version}` format of requirements.txt that pip uses. impl std::fmt::Display for DisplayResolutionGraph<'_> { fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + // Reduce the graph, such that all nodes for a single package are combined, regardless of + // the extras. + // + // For example, `flask` and `flask[dotenv]` should be reduced into a single `flask[dotenv]` + // node. + let petgraph = { + let mut petgraph = + petgraph::graph::Graph::::with_capacity( + self.resolution.petgraph.node_count(), + self.resolution.petgraph.edge_count(), + ); + let mut inverse = FxHashMap::with_capacity_and_hasher( + self.resolution.petgraph.node_count(), + BuildHasherDefault::default(), + ); + + // Re-add the nodes to the reduced graph. + for index in self.resolution.petgraph.node_indices() { + let dist = &self.resolution.petgraph[index]; + + if let Some(index) = inverse.get(dist.name()) { + let node: &mut AnnotatedDist = &mut petgraph[*index]; + node.extras.extend(dist.extras.iter().cloned()); + node.extras.sort_unstable(); + node.extras.dedup(); + } else { + let index = petgraph.add_node(dist.clone()); + inverse.insert(dist.name(), index); + } + } + + // Re-add the edges to the reduced graph. + for edge in self.resolution.petgraph.edge_indices() { + let (source, target) = self.resolution.petgraph.edge_endpoints(edge).unwrap(); + let source = inverse[self.resolution.petgraph[source].name()]; + let target = inverse[self.resolution.petgraph[target].name()]; + petgraph.update_edge(source, target, ()); + } + + petgraph + }; + // Collect all packages. - let mut nodes = self - .resolution - .petgraph + let mut nodes = petgraph .node_indices() .filter_map(|index| { - let dist = &self.resolution.petgraph[index]; + let dist = &petgraph[index]; let name = dist.name(); if self.no_emit_packages.contains(name) { return None; @@ -119,11 +162,9 @@ impl std::fmt::Display for DisplayResolutionGraph<'_> { // package (e.g., `# via mypy`). if self.include_annotations { // Display all dependencies. - let mut edges = self - .resolution - .petgraph + let mut edges = petgraph .edges_directed(index, Direction::Incoming) - .map(|edge| &self.resolution.petgraph[edge.source()]) + .map(|edge| &petgraph[edge.source()]) .collect::>(); edges.sort_unstable_by_key(|package| package.name()); diff --git a/crates/uv-resolver/src/resolution/graph.rs b/crates/uv-resolver/src/resolution/graph.rs index 2aedc0e6a..70478daa7 100644 --- a/crates/uv-resolver/src/resolution/graph.rs +++ b/crates/uv-resolver/src/resolution/graph.rs @@ -46,102 +46,17 @@ impl ResolutionGraph { state: &State, preferences: &Preferences, ) -> anyhow::Result { - // Collect and validate the extras. - let mut extras = FxHashMap::default(); - let mut diagnostics = Vec::new(); - for (package, version) in selection { - match &**package { - PubGrubPackageInner::Package { - name, - extra: Some(extra), - marker: None, - url: None, - } => { - let dist = PubGrubDistribution::from_registry(name, version); - - let response = distributions.get(&dist.version_id()).unwrap_or_else(|| { - panic!( - "Every package should have metadata: {:?}", - dist.version_id() - ) - }); - - let MetadataResponse::Found(archive) = &*response else { - panic!( - "Every package should have metadata: {:?}", - dist.version_id() - ) - }; - - if archive.metadata.provides_extras.contains(extra) { - extras - .entry(name.clone()) - .or_insert_with(Vec::new) - .push(extra.clone()); - } else { - let dist = pins - .get(name, version) - .unwrap_or_else(|| panic!("Every package should be pinned: {name:?}")) - .clone(); - - diagnostics.push(ResolutionDiagnostic::MissingExtra { - dist, - extra: extra.clone(), - }); - } - } - PubGrubPackageInner::Package { - name, - extra: Some(extra), - marker: None, - url: Some(url), - } => { - let dist = PubGrubDistribution::from_url(name, url); - - let response = distributions.get(&dist.version_id()).unwrap_or_else(|| { - panic!( - "Every package should have metadata: {:?}", - dist.version_id() - ) - }); - - let MetadataResponse::Found(archive) = &*response else { - panic!( - "Every package should have metadata: {:?}", - dist.version_id() - ) - }; - - if archive.metadata.provides_extras.contains(extra) { - extras - .entry(name.clone()) - .or_insert_with(Vec::new) - .push(extra.clone()); - } else { - let dist = Dist::from_url(name.clone(), url_to_precise(url.clone()))?; - - diagnostics.push(ResolutionDiagnostic::MissingExtra { - dist: dist.into(), - extra: extra.clone(), - }); - } - } - _ => {} - }; - } - // Add every package to the graph. - // TODO(charlie): petgraph is a really heavy and unnecessary dependency here. We should - // write our own graph, given that our requirements are so simple. let mut petgraph = petgraph::graph::Graph::with_capacity(selection.len(), selection.len()); let mut inverse = FxHashMap::with_capacity_and_hasher(selection.len(), BuildHasherDefault::default()); + let mut diagnostics = Vec::new(); for (package, version) in selection { match &**package { PubGrubPackageInner::Package { name, - extra: None, + extra, marker: None, url: None, } => { @@ -213,8 +128,18 @@ impl ResolutionGraph { archive.metadata.clone() }; + // Validate the extra. + if let Some(extra) = extra { + if !metadata.provides_extras.contains(extra) { + diagnostics.push(ResolutionDiagnostic::MissingExtra { + dist: dist.clone(), + extra: extra.clone(), + }); + } + } + // Extract the extras. - let extras = extras.get(name).cloned().unwrap_or_default(); + let extras = extra.clone().map(|extra| vec![extra]).unwrap_or_default(); // Add the distribution to the graph. let index = petgraph.add_node(AnnotatedDist { @@ -223,16 +148,16 @@ impl ResolutionGraph { hashes, metadata, }); - inverse.insert(name, index); + inverse.insert((name, extra), index); } + PubGrubPackageInner::Package { name, - extra: None, + extra, marker: None, url: Some(url), } => { // Create the distribution. - let dist = Dist::from_url(name.clone(), url_to_precise(url.clone()))?; // Extract the hashes, preserving those that were already present in the @@ -275,8 +200,18 @@ impl ResolutionGraph { archive.metadata.clone() }; + // Validate the extra. + if let Some(extra) = extra { + if !metadata.provides_extras.contains(extra) { + diagnostics.push(ResolutionDiagnostic::MissingExtra { + dist: dist.clone().into(), + extra: extra.clone(), + }); + } + } + // Extract the extras. - let extras = extras.get(name).cloned().unwrap_or_default(); + let extras = extra.clone().map(|extra| vec![extra]).unwrap_or_default(); // Add the distribution to the graph. let index = petgraph.add_node(AnnotatedDist { @@ -285,8 +220,9 @@ impl ResolutionGraph { hashes, metadata, }); - inverse.insert(name, index); + inverse.insert((name, extra), index); } + _ => {} }; } @@ -305,14 +241,22 @@ impl ResolutionGraph { continue; } + if !self_version.contains(version) { + continue; + } + let PubGrubPackageInner::Package { - name: self_name, .. + name: self_name, + extra: self_extra, + .. } = &**self_package else { continue; }; + let PubGrubPackageInner::Package { name: dependency_name, + extra: dependency_extra, .. } = &**dependency_package else { @@ -324,11 +268,9 @@ impl ResolutionGraph { continue; } - if self_version.contains(version) { - let self_index = &inverse[self_name]; - let dependency_index = &inverse[dependency_name]; - petgraph.update_edge(*self_index, *dependency_index, ()); - } + let self_index = &inverse[&(self_name, self_extra)]; + let dependency_index = &inverse[&(dependency_name, dependency_extra)]; + petgraph.update_edge(*self_index, *dependency_index, ()); } } } @@ -339,14 +281,18 @@ impl ResolutionGraph { }) } - /// Return the number of packages in the graph. + /// Return the number of distinct packages in the graph. pub fn len(&self) -> usize { - self.petgraph.node_count() + self.petgraph + .node_indices() + .map(|index| &self.petgraph[index]) + .filter(|dist| dist.extras.is_empty()) + .count() } /// Return `true` if there are no packages in the graph. pub fn is_empty(&self) -> bool { - self.petgraph.node_count() == 0 + self.len() == 0 } /// Returns `true` if the graph contains the given package.