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.
This commit is contained in:
Charlie Marsh 2024-03-19 19:59:32 -04:00 committed by GitHub
parent ad396a7cff
commit ab99a18cbc
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
5 changed files with 259 additions and 54 deletions

View file

@ -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<Dist, Range<Version>, petgraph::Directed>,
/// The metadata for every distribution in this resolution.
hashes: FxHashMap<PackageName, Vec<Hashes>>,
/// The enabled extras for every distribution in this resolution.
extras: FxHashMap<PackageName, Vec<ExtraName>>,
/// 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;

View file

@ -54,6 +54,7 @@ pub(crate) async fn pip_compile(
upgrade: Upgrade,
generate_hashes: bool,
no_emit_packages: Vec<PackageName>,
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,
)

View file

@ -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<String>,
}
@ -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() {

View file

@ -316,6 +316,14 @@ struct PipCompileArgs {
#[clap(long, short)]
output_file: Option<PathBuf>,
/// 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<ExitStatus> {
upgrade,
args.generate_hashes,
args.no_emit_package,
args.no_strip_extras,
!args.no_annotate,
!args.no_header,
args.emit_index_url,

View file

@ -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]