Support multiple extras in universal pip compile output (#8960)

## 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 https://github.com/astral-sh/uv/issues/8915.
This commit is contained in:
Charlie Marsh
2024-11-08 18:22:29 -05:00
committed by GitHub
parent f399a52719
commit 75c26229a2
3 changed files with 479 additions and 52 deletions
+79 -50
View File
@@ -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);
}
}
}
@@ -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,8 +181,10 @@ 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.
@@ -187,6 +192,7 @@ pub(crate) enum RequirementsTxtComparator<'a> {
name: &'a PackageName,
version: &'a Version,
url: Option<Cow<'a, str>>,
extras: &'a [ExtraName],
},
}
+392
View File
@@ -8560,6 +8560,398 @@ fn universal_marker_propagation() -> Result<()> {
Ok(())
}
#[test]
fn universal_disjoint_extra() -> Result<()> {
let context = TestContext::new("3.12");
let requirements_in = context.temp_dir.child("requirements.in");
requirements_in.write_str(indoc::indoc! {r"
flask[async]; sys_platform == 'linux'
flask[dotenv]; sys_platform == 'darwin'
"})?;
uv_snapshot!(context.filters(), windows_filters=false, context.pip_compile()
.arg("requirements.in")
.arg("--universal"), @r###"
success: true
exit_code: 0
----- stdout -----
# This file was autogenerated by uv via the following command:
# uv pip compile --cache-dir [CACHE_DIR] requirements.in --universal
asgiref==3.8.1 ; sys_platform == 'linux'
# via flask
blinker==1.7.0 ; sys_platform == 'darwin' or sys_platform == 'linux'
# via flask
click==8.1.7 ; sys_platform == 'darwin' or sys_platform == 'linux'
# via flask
colorama==0.4.6 ; (platform_system == 'Windows' and sys_platform == 'darwin') or (platform_system == 'Windows' and sys_platform == 'linux')
# via click
flask==3.0.2 ; sys_platform == 'linux'
# via -r requirements.in
itsdangerous==2.1.2 ; sys_platform == 'darwin' or sys_platform == 'linux'
# via flask
jinja2==3.1.3 ; sys_platform == 'darwin' or sys_platform == 'linux'
# via flask
markupsafe==2.1.5 ; sys_platform == 'darwin' or sys_platform == 'linux'
# via
# jinja2
# werkzeug
python-dotenv==1.0.1 ; sys_platform == 'darwin'
# via flask
werkzeug==3.0.1 ; sys_platform == 'darwin' or sys_platform == 'linux'
# via flask
----- stderr -----
Resolved 10 packages in [TIME]
"###
);
Ok(())
}
#[test]
fn universal_disjoint_extra_no_strip() -> Result<()> {
let context = TestContext::new("3.12");
let requirements_in = context.temp_dir.child("requirements.in");
requirements_in.write_str(indoc::indoc! {r"
flask[async]; sys_platform == 'linux'
flask[dotenv]; sys_platform == 'darwin'
"})?;
uv_snapshot!(context.filters(), windows_filters=false, context.pip_compile()
.arg("requirements.in")
.arg("--universal")
.arg("--no-strip-extras"), @r###"
success: true
exit_code: 0
----- stdout -----
# This file was autogenerated by uv via the following command:
# uv pip compile --cache-dir [CACHE_DIR] requirements.in --universal --no-strip-extras
asgiref==3.8.1 ; sys_platform == 'linux'
# via flask
blinker==1.7.0 ; sys_platform == 'darwin' or sys_platform == 'linux'
# via flask
click==8.1.7 ; sys_platform == 'darwin' or sys_platform == 'linux'
# via flask
colorama==0.4.6 ; (platform_system == 'Windows' and sys_platform == 'darwin') or (platform_system == 'Windows' and sys_platform == 'linux')
# via click
flask==3.0.2 ; sys_platform == 'darwin' or sys_platform == 'linux'
# via -r requirements.in
flask[async]==3.0.2 ; sys_platform == 'linux'
# via -r requirements.in
flask[dotenv]==3.0.2 ; sys_platform == 'darwin'
# via -r requirements.in
itsdangerous==2.1.2 ; sys_platform == 'darwin' or sys_platform == 'linux'
# via flask
jinja2==3.1.3 ; sys_platform == 'darwin' or sys_platform == 'linux'
# via flask
markupsafe==2.1.5 ; sys_platform == 'darwin' or sys_platform == 'linux'
# via
# jinja2
# werkzeug
python-dotenv==1.0.1 ; sys_platform == 'darwin'
# via flask
werkzeug==3.0.1 ; sys_platform == 'darwin' or sys_platform == 'linux'
# via flask
----- stderr -----
Resolved 10 packages in [TIME]
"###
);
Ok(())
}
#[test]
fn universal_overlap_extra_base() -> Result<()> {
let context = TestContext::new("3.12");
let requirements_in = context.temp_dir.child("requirements.in");
requirements_in.write_str(indoc::indoc! {r"
flask
flask[dotenv]; sys_platform == 'darwin'
"})?;
uv_snapshot!(context.filters(), windows_filters=false, context.pip_compile()
.arg("requirements.in")
.arg("--universal"), @r###"
success: true
exit_code: 0
----- stdout -----
# This file was autogenerated by uv via the following command:
# uv pip compile --cache-dir [CACHE_DIR] requirements.in --universal
blinker==1.7.0
# via flask
click==8.1.7
# via flask
colorama==0.4.6 ; platform_system == 'Windows'
# via click
flask==3.0.2
# via -r requirements.in
itsdangerous==2.1.2
# via flask
jinja2==3.1.3
# via flask
markupsafe==2.1.5
# via
# jinja2
# werkzeug
python-dotenv==1.0.1 ; sys_platform == 'darwin'
# via flask
werkzeug==3.0.1
# via flask
----- stderr -----
Resolved 9 packages in [TIME]
"###
);
Ok(())
}
#[test]
fn universal_overlap_extra_base_no_strip() -> Result<()> {
let context = TestContext::new("3.12");
let requirements_in = context.temp_dir.child("requirements.in");
requirements_in.write_str(indoc::indoc! {r"
flask
flask[dotenv]; sys_platform == 'darwin'
"})?;
uv_snapshot!(context.filters(), windows_filters=false, context.pip_compile()
.arg("requirements.in")
.arg("--universal")
.arg("--no-strip-extras"), @r###"
success: true
exit_code: 0
----- stdout -----
# This file was autogenerated by uv via the following command:
# uv pip compile --cache-dir [CACHE_DIR] requirements.in --universal --no-strip-extras
blinker==1.7.0
# via flask
click==8.1.7
# via flask
colorama==0.4.6 ; platform_system == 'Windows'
# via click
flask==3.0.2
# via -r requirements.in
flask[dotenv]==3.0.2 ; sys_platform == 'darwin'
# via -r requirements.in
itsdangerous==2.1.2
# via flask
jinja2==3.1.3
# via flask
markupsafe==2.1.5
# via
# jinja2
# werkzeug
python-dotenv==1.0.1 ; sys_platform == 'darwin'
# via flask
werkzeug==3.0.1
# via flask
----- stderr -----
Resolved 9 packages in [TIME]
"###
);
Ok(())
}
#[test]
fn universal_overlap_extras() -> Result<()> {
let context = TestContext::new("3.12");
let requirements_in = context.temp_dir.child("requirements.in");
requirements_in.write_str(indoc::indoc! {r"
flask[async]; sys_platform == 'linux' or sys_platform == 'darwin'
flask[dotenv]; sys_platform == 'darwin'
"})?;
uv_snapshot!(context.filters(), windows_filters=false, context.pip_compile()
.arg("requirements.in")
.arg("--universal"), @r###"
success: true
exit_code: 0
----- stdout -----
# This file was autogenerated by uv via the following command:
# uv pip compile --cache-dir [CACHE_DIR] requirements.in --universal
asgiref==3.8.1 ; sys_platform == 'darwin' or sys_platform == 'linux'
# via flask
blinker==1.7.0 ; sys_platform == 'darwin' or sys_platform == 'linux'
# via flask
click==8.1.7 ; sys_platform == 'darwin' or sys_platform == 'linux'
# via flask
colorama==0.4.6 ; (platform_system == 'Windows' and sys_platform == 'darwin') or (platform_system == 'Windows' and sys_platform == 'linux')
# via click
flask==3.0.2 ; sys_platform == 'darwin' or sys_platform == 'linux'
# via -r requirements.in
itsdangerous==2.1.2 ; sys_platform == 'darwin' or sys_platform == 'linux'
# via flask
jinja2==3.1.3 ; sys_platform == 'darwin' or sys_platform == 'linux'
# via flask
markupsafe==2.1.5 ; sys_platform == 'darwin' or sys_platform == 'linux'
# via
# jinja2
# werkzeug
python-dotenv==1.0.1 ; sys_platform == 'darwin'
# via flask
werkzeug==3.0.1 ; sys_platform == 'darwin' or sys_platform == 'linux'
# via flask
----- stderr -----
Resolved 10 packages in [TIME]
"###
);
Ok(())
}
#[test]
fn universal_overlap_extras_no_strip() -> Result<()> {
let context = TestContext::new("3.12");
let requirements_in = context.temp_dir.child("requirements.in");
requirements_in.write_str(indoc::indoc! {r"
flask[async]; sys_platform == 'linux' or sys_platform == 'darwin'
flask[dotenv]; sys_platform == 'darwin'
"})?;
uv_snapshot!(context.filters(), windows_filters=false, context.pip_compile()
.arg("requirements.in")
.arg("--universal")
.arg("--no-strip-extras"), @r###"
success: true
exit_code: 0
----- stdout -----
# This file was autogenerated by uv via the following command:
# uv pip compile --cache-dir [CACHE_DIR] requirements.in --universal --no-strip-extras
asgiref==3.8.1 ; sys_platform == 'darwin' or sys_platform == 'linux'
# via flask
blinker==1.7.0 ; sys_platform == 'darwin' or sys_platform == 'linux'
# via flask
click==8.1.7 ; sys_platform == 'darwin' or sys_platform == 'linux'
# via flask
colorama==0.4.6 ; (platform_system == 'Windows' and sys_platform == 'darwin') or (platform_system == 'Windows' and sys_platform == 'linux')
# via click
flask[async]==3.0.2 ; sys_platform == 'darwin' or sys_platform == 'linux'
# via -r requirements.in
flask[dotenv]==3.0.2 ; sys_platform == 'darwin'
# via -r requirements.in
itsdangerous==2.1.2 ; sys_platform == 'darwin' or sys_platform == 'linux'
# via flask
jinja2==3.1.3 ; sys_platform == 'darwin' or sys_platform == 'linux'
# via flask
markupsafe==2.1.5 ; sys_platform == 'darwin' or sys_platform == 'linux'
# via
# jinja2
# werkzeug
python-dotenv==1.0.1 ; sys_platform == 'darwin'
# via flask
werkzeug==3.0.1 ; sys_platform == 'darwin' or sys_platform == 'linux'
# via flask
----- stderr -----
Resolved 10 packages in [TIME]
"###
);
Ok(())
}
#[test]
fn universal_identical_extras() -> Result<()> {
let context = TestContext::new("3.12");
let requirements_in = context.temp_dir.child("requirements.in");
requirements_in.write_str(indoc::indoc! {r"
flask[async]; sys_platform == 'darwin'
flask[dotenv]; sys_platform == 'darwin'
"})?;
uv_snapshot!(context.filters(), windows_filters=false, context.pip_compile()
.arg("requirements.in")
.arg("--universal"), @r###"
success: true
exit_code: 0
----- stdout -----
# This file was autogenerated by uv via the following command:
# uv pip compile --cache-dir [CACHE_DIR] requirements.in --universal
asgiref==3.8.1 ; sys_platform == 'darwin'
# via flask
blinker==1.7.0 ; sys_platform == 'darwin'
# via flask
click==8.1.7 ; sys_platform == 'darwin'
# via flask
colorama==0.4.6 ; platform_system == 'Windows' and sys_platform == 'darwin'
# via click
flask==3.0.2 ; sys_platform == 'darwin'
# via -r requirements.in
itsdangerous==2.1.2 ; sys_platform == 'darwin'
# via flask
jinja2==3.1.3 ; sys_platform == 'darwin'
# via flask
markupsafe==2.1.5 ; sys_platform == 'darwin'
# via
# jinja2
# werkzeug
python-dotenv==1.0.1 ; sys_platform == 'darwin'
# via flask
werkzeug==3.0.1 ; sys_platform == 'darwin'
# via flask
----- stderr -----
Resolved 10 packages in [TIME]
"###
);
Ok(())
}
#[test]
fn universal_identical_extras_no_strip() -> Result<()> {
let context = TestContext::new("3.12");
let requirements_in = context.temp_dir.child("requirements.in");
requirements_in.write_str(indoc::indoc! {r"
flask[async]; sys_platform == 'darwin'
flask[dotenv]; sys_platform == 'darwin'
"})?;
uv_snapshot!(context.filters(), windows_filters=false, context.pip_compile()
.arg("requirements.in")
.arg("--universal")
.arg("--no-strip-extras"), @r###"
success: true
exit_code: 0
----- stdout -----
# This file was autogenerated by uv via the following command:
# uv pip compile --cache-dir [CACHE_DIR] requirements.in --universal --no-strip-extras
asgiref==3.8.1 ; sys_platform == 'darwin'
# via flask
blinker==1.7.0 ; sys_platform == 'darwin'
# via flask
click==8.1.7 ; sys_platform == 'darwin'
# via flask
colorama==0.4.6 ; platform_system == 'Windows' and sys_platform == 'darwin'
# via click
flask[async, dotenv]==3.0.2 ; sys_platform == 'darwin'
# via -r requirements.in
itsdangerous==2.1.2 ; sys_platform == 'darwin'
# via flask
jinja2==3.1.3 ; sys_platform == 'darwin'
# via flask
markupsafe==2.1.5 ; sys_platform == 'darwin'
# via
# jinja2
# werkzeug
python-dotenv==1.0.1 ; sys_platform == 'darwin'
# via flask
werkzeug==3.0.1 ; sys_platform == 'darwin'
# via flask
----- stderr -----
Resolved 10 packages in [TIME]
"###
);
Ok(())
}
/// Resolve a package from a `requirements.in` file, with a `constraints.txt` file pinning one of
/// its transitive dependencies to a specific version.
#[test]