mirror of
https://github.com/astral-sh/uv.git
synced 2025-10-02 23:04:37 +00:00
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.
This commit is contained in:
Charlie Marsh
GitHub
parent
0edb660205
commit
19c91e7dac
2 changed files with 97 additions and 110 deletions
|
|
@ -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::<AnnotatedDist, (), petgraph::Directed>::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::<Vec<_>>();
|
||||
edges.sort_unstable_by_key(|package| package.name());
|
||||
|
||||
|
|
|
|||
|
|
@ -46,102 +46,17 @@ impl ResolutionGraph {
|
|||
state: &State<UvDependencyProvider>,
|
||||
preferences: &Preferences,
|
||||
) -> anyhow::Result<Self, ResolveError> {
|
||||
// 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.
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue