From ab99a18cbca41c7364d8c30e9f90369e7a798691 Mon Sep 17 00:00:00 2001 From: Charlie Marsh Date: Tue, 19 Mar 2024 19:59:32 -0400 Subject: [PATCH] Implement `--no-strip-extras` to preserve extras in compilation (#2555) ## Summary We strip extras by default, but there are some valid use-cases in which they're required (see the linked issue). This PR doesn't change our default, but it does add `--no-strip-extras`, which lets users preserve extras in the output requirements. Closes https://github.com/astral-sh/uv/issues/1595. --- crates/uv-resolver/src/resolution.rs | 150 ++++++++++++++++++-------- crates/uv/src/commands/pip_compile.rs | 2 + crates/uv/src/compat/mod.rs | 13 +-- crates/uv/src/main.rs | 9 ++ crates/uv/tests/pip_compile.rs | 139 ++++++++++++++++++++++++ 5 files changed, 259 insertions(+), 54 deletions(-) diff --git a/crates/uv-resolver/src/resolution.rs b/crates/uv-resolver/src/resolution.rs index e8a2ef4be..909b0d340 100644 --- a/crates/uv-resolver/src/resolution.rs +++ b/crates/uv-resolver/src/resolution.rs @@ -3,6 +3,7 @@ use std::hash::BuildHasherDefault; use anyhow::Result; use dashmap::DashMap; +use itertools::Itertools; use owo_colors::OwoColorize; use petgraph::visit::EdgeRef; use petgraph::Direction; @@ -46,6 +47,8 @@ pub struct ResolutionGraph { petgraph: petgraph::graph::Graph, petgraph::Directed>, /// The metadata for every distribution in this resolution. hashes: FxHashMap>, + /// The enabled extras for every distribution in this resolution. + extras: FxHashMap>, /// The set of editable requirements in this resolution. editables: Editables, /// Any diagnostics that were encountered while building the graph. @@ -70,6 +73,7 @@ impl ResolutionGraph { let mut petgraph = petgraph::graph::Graph::with_capacity(selection.len(), selection.len()); let mut hashes = FxHashMap::with_capacity_and_hasher(selection.len(), BuildHasherDefault::default()); + let mut extras = FxHashMap::default(); let mut diagnostics = Vec::new(); // Add every package to the graph. @@ -140,7 +144,12 @@ impl ResolutionGraph { let dist = PubGrubDistribution::from_registry(package_name, version); if let Some((editable, metadata)) = editables.get(package_name) { - if !metadata.provides_extras.contains(extra) { + if metadata.provides_extras.contains(extra) { + extras + .entry(package_name.clone()) + .or_insert_with(Vec::new) + .push(extra.clone()); + } else { let pinned_package = Dist::from_editable(package_name.clone(), editable.clone())?; @@ -157,7 +166,12 @@ impl ResolutionGraph { ) }); - if !metadata.provides_extras.contains(extra) { + if metadata.provides_extras.contains(extra) { + extras + .entry(package_name.clone()) + .or_insert_with(Vec::new) + .push(extra.clone()); + } else { let pinned_package = pins .get(package_name, version) .unwrap_or_else(|| { @@ -177,7 +191,12 @@ impl ResolutionGraph { let dist = PubGrubDistribution::from_url(package_name, url); if let Some((editable, metadata)) = editables.get(package_name) { - if !metadata.provides_extras.contains(extra) { + if metadata.provides_extras.contains(extra) { + extras + .entry(package_name.clone()) + .or_insert_with(Vec::new) + .push(extra.clone()); + } else { let pinned_package = Dist::from_editable(package_name.clone(), editable.clone())?; @@ -194,7 +213,12 @@ impl ResolutionGraph { ) }); - if !metadata.provides_extras.contains(extra) { + if metadata.provides_extras.contains(extra) { + extras + .entry(package_name.clone()) + .or_insert_with(Vec::new) + .push(extra.clone()); + } else { let url = redirects.get(url).map_or_else( || url.clone(), |precise| apply_redirect(url, precise.value()), @@ -259,6 +283,7 @@ impl ResolutionGraph { Ok(Self { petgraph, hashes, + extras, editables, diagnostics, }) @@ -301,6 +326,8 @@ pub struct DisplayResolutionGraph<'a> { no_emit_packages: &'a [PackageName], /// Whether to include hashes in the output. show_hashes: bool, + /// Whether to include extras in the output (e.g., `black[colorama]`). + include_extras: bool, /// Whether to include annotations in the output, to indicate which dependency or dependencies /// requested each package. include_annotations: bool, @@ -311,7 +338,14 @@ pub struct DisplayResolutionGraph<'a> { impl<'a> From<&'a ResolutionGraph> for DisplayResolutionGraph<'a> { fn from(resolution: &'a ResolutionGraph) -> Self { - Self::new(resolution, &[], false, true, AnnotationStyle::default()) + Self::new( + resolution, + &[], + false, + false, + true, + AnnotationStyle::default(), + ) } } @@ -321,6 +355,7 @@ impl<'a> DisplayResolutionGraph<'a> { underlying: &'a ResolutionGraph, no_emit_packages: &'a [PackageName], show_hashes: bool, + include_extras: bool, include_annotations: bool, annotation_style: AnnotationStyle, ) -> DisplayResolutionGraph<'a> { @@ -328,49 +363,70 @@ impl<'a> DisplayResolutionGraph<'a> { resolution: underlying, no_emit_packages, show_hashes, + include_extras, include_annotations, annotation_style, } } } +#[derive(Debug)] +enum Node<'a> { + /// A node linked to an editable distribution. + Editable(&'a PackageName, &'a LocalEditable), + /// A node linked to a non-editable distribution. + Distribution(&'a PackageName, &'a Dist, &'a [ExtraName]), +} + +#[derive(Debug, PartialEq, Eq, PartialOrd, Ord)] +enum NodeKey<'a> { + /// A node linked to an editable distribution, sorted by verbatim representation. + Editable(Cow<'a, str>), + /// A node linked to a non-editable distribution, sorted by package name. + Distribution(&'a PackageName), +} + +impl<'a> Node<'a> { + /// Return the name of the package. + fn name(&self) -> &'a PackageName { + match self { + Node::Editable(name, _) => name, + Node::Distribution(name, _, _) => name, + } + } + + /// Return a comparable key for the node. + fn key(&self) -> NodeKey<'a> { + match self { + Node::Editable(_, editable) => NodeKey::Editable(editable.verbatim()), + Node::Distribution(name, _, _) => NodeKey::Distribution(name), + } + } +} + +impl Verbatim for Node<'_> { + fn verbatim(&self) -> Cow<'_, str> { + match self { + Node::Editable(_, editable) => Cow::Owned(format!("-e {}", editable.verbatim())), + Node::Distribution(_, dist, &[]) => dist.verbatim(), + Node::Distribution(_, dist, extras) => { + let mut extras = extras.to_vec(); + extras.sort_unstable(); + extras.dedup(); + Cow::Owned(format!( + "{}[{}]{}", + dist.name(), + extras.into_iter().join(", "), + dist.version_or_url().verbatim() + )) + } + } + } +} + /// 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 { - #[derive(Debug)] - enum Node<'a> { - /// A node linked to an editable distribution. - Editable(&'a PackageName, &'a LocalEditable), - /// A node linked to a non-editable distribution. - Distribution(&'a PackageName, &'a Dist), - } - - #[derive(Debug, PartialEq, Eq, PartialOrd, Ord)] - enum NodeKey<'a> { - /// A node linked to an editable distribution, sorted by verbatim representation. - Editable(Cow<'a, str>), - /// A node linked to a non-editable distribution, sorted by package name. - Distribution(&'a PackageName), - } - - impl<'a> Node<'a> { - /// Return the name of the package. - fn name(&self) -> &'a PackageName { - match self { - Node::Editable(name, _) => name, - Node::Distribution(name, _) => name, - } - } - - /// Return a comparable key for the node. - fn key(&self) -> NodeKey<'a> { - match self { - Node::Editable(_, editable) => NodeKey::Editable(editable.verbatim()), - Node::Distribution(name, _) => NodeKey::Distribution(name), - } - } - } - // Collect all packages. let mut nodes = self .resolution @@ -385,8 +441,17 @@ impl std::fmt::Display for DisplayResolutionGraph<'_> { let node = if let Some((editable, _)) = self.resolution.editables.get(name) { Node::Editable(name, editable) + } else if self.include_extras { + Node::Distribution( + name, + dist, + self.resolution + .extras + .get(name) + .map_or(&[], |extras| extras.as_slice()), + ) } else { - Node::Distribution(name, dist) + Node::Distribution(name, dist, &[]) }; Some((index, node)) }) @@ -398,10 +463,7 @@ impl std::fmt::Display for DisplayResolutionGraph<'_> { // Print out the dependency graph. for (index, node) in nodes { // Display the node itself. - let mut line = match node { - Node::Distribution(_, dist) => format!("{}", dist.verbatim()), - Node::Editable(_, editable) => format!("-e {}", editable.verbatim()), - }; + let mut line = node.verbatim().to_string(); // Display the distribution hashes, if any. let mut has_hashes = false; diff --git a/crates/uv/src/commands/pip_compile.rs b/crates/uv/src/commands/pip_compile.rs index 8afd94c85..aae20fae1 100644 --- a/crates/uv/src/commands/pip_compile.rs +++ b/crates/uv/src/commands/pip_compile.rs @@ -54,6 +54,7 @@ pub(crate) async fn pip_compile( upgrade: Upgrade, generate_hashes: bool, no_emit_packages: Vec, + include_extras: bool, include_annotations: bool, include_header: bool, include_index_url: bool, @@ -408,6 +409,7 @@ pub(crate) async fn pip_compile( &resolution, &no_emit_packages, generate_hashes, + include_extras, include_annotations, annotation_style, ) diff --git a/crates/uv/src/compat/mod.rs b/crates/uv/src/compat/mod.rs index bcea487a7..dc7ab0d57 100644 --- a/crates/uv/src/compat/mod.rs +++ b/crates/uv/src/compat/mod.rs @@ -72,9 +72,6 @@ pub(crate) struct PipCompileCompatArgs { #[clap(long, hide = true)] strip_extras: bool, - #[clap(long, hide = true)] - no_strip_extras: bool, - #[clap(long, hide = true)] pip_args: Option, } @@ -194,13 +191,9 @@ impl CompatArgs for PipCompileCompatArgs { } if self.strip_extras { - warn_user!("pip-compile's `--strip-extras` has no effect (uv always strips extras)."); - } - - if self.no_strip_extras { - return Err(anyhow!( - "pip-compile's `--no-strip-extras` is unsupported (uv always strips extras)." - )); + warn_user!( + "pip-compile's `--strip-extras` has no effect (uv strips extras by default)." + ); } if self.pip_args.is_some() { diff --git a/crates/uv/src/main.rs b/crates/uv/src/main.rs index ee99de892..bf7429fd4 100644 --- a/crates/uv/src/main.rs +++ b/crates/uv/src/main.rs @@ -316,6 +316,14 @@ struct PipCompileArgs { #[clap(long, short)] output_file: Option, + /// Include extras in the output file. + /// + /// By default, `uv` strips extras, as any packages pulled in by the extras are already included + /// as dependencies in the output file directly. Further, output files generated with + /// `--no-strip-extras` cannot be used as constraints files in `install` and `sync` invocations. + #[clap(long)] + no_strip_extras: bool, + /// Exclude comment annotations indicating the source of each package. #[clap(long)] no_annotate: bool, @@ -1505,6 +1513,7 @@ async fn run() -> Result { upgrade, args.generate_hashes, args.no_emit_package, + args.no_strip_extras, !args.no_annotate, !args.no_header, args.emit_index_url, diff --git a/crates/uv/tests/pip_compile.rs b/crates/uv/tests/pip_compile.rs index 9430c6f1f..8c1747bc0 100644 --- a/crates/uv/tests/pip_compile.rs +++ b/crates/uv/tests/pip_compile.rs @@ -4175,6 +4175,145 @@ fn editable_invalid_extra() -> Result<()> { Ok(()) } +/// Resolve a package with `--no-strip-extras`. +#[test] +fn no_strip_extra() -> Result<()> { + let context = TestContext::new("3.12"); + let requirements_in = context.temp_dir.child("requirements.in"); + requirements_in.write_str("flask[dotenv]")?; + + uv_snapshot!(context.compile() + .arg("requirements.in") + .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] --exclude-newer 2023-11-18T12:00:00Z requirements.in --no-strip-extras + blinker==1.7.0 + # via flask + click==8.1.7 + # via flask + flask[dotenv]==3.0.0 + itsdangerous==2.1.2 + # via flask + jinja2==3.1.2 + # via flask + markupsafe==2.1.3 + # via + # jinja2 + # werkzeug + python-dotenv==1.0.0 + # via flask + werkzeug==3.0.1 + # via flask + + ----- stderr ----- + Resolved 8 packages in [TIME] + "### + ); + + Ok(()) +} + +/// Resolve a package with `--no-strip-extras`. +#[test] +#[cfg(not(windows))] +fn no_strip_extras() -> Result<()> { + let context = TestContext::new("3.12"); + let requirements_in = context.temp_dir.child("requirements.in"); + requirements_in.write_str("anyio[trio]\nanyio[doc]")?; + + uv_snapshot!(context.compile() + .arg("requirements.in") + .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] --exclude-newer 2023-11-18T12:00:00Z requirements.in --no-strip-extras + alabaster==0.7.13 + # via sphinx + anyio[doc, trio]==4.0.0 + attrs==23.1.0 + # via + # outcome + # trio + babel==2.13.1 + # via sphinx + certifi==2023.11.17 + # via requests + charset-normalizer==3.3.2 + # via requests + docutils==0.20.1 + # via sphinx + idna==3.4 + # via + # anyio + # requests + # trio + imagesize==1.4.1 + # via sphinx + jinja2==3.1.2 + # via sphinx + markupsafe==2.1.3 + # via jinja2 + outcome==1.3.0.post0 + # via trio + packaging==23.2 + # via + # anyio + # sphinx + pygments==2.16.1 + # via sphinx + requests==2.31.0 + # via sphinx + setuptools==68.2.2 + # via babel + sniffio==1.3.0 + # via + # anyio + # trio + snowballstemmer==2.2.0 + # via sphinx + sortedcontainers==2.4.0 + # via trio + sphinx==7.2.6 + # via + # anyio + # sphinx-autodoc-typehints + # sphinxcontrib-applehelp + # sphinxcontrib-devhelp + # sphinxcontrib-htmlhelp + # sphinxcontrib-qthelp + # sphinxcontrib-serializinghtml + sphinx-autodoc-typehints==1.25.2 + # via anyio + sphinxcontrib-applehelp==1.0.7 + # via sphinx + sphinxcontrib-devhelp==1.0.5 + # via sphinx + sphinxcontrib-htmlhelp==2.0.4 + # via sphinx + sphinxcontrib-jsmath==1.0.1 + # via sphinx + sphinxcontrib-qthelp==1.0.6 + # via sphinx + sphinxcontrib-serializinghtml==1.1.9 + # via sphinx + trio==0.23.1 + # via anyio + urllib3==2.1.0 + # via requests + + ----- stderr ----- + Resolved 29 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]