Skip to content

Commit

Permalink
Implement --no-strip-extras to preserve extras in compilation (#2555)
Browse files Browse the repository at this point in the history 
## 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 #1595.
  • Loading branch information
@charliermarsh
charliermarsh committed Mar 19, 2024
1 parent ad396a7 commit ab99a18
Show file tree
Hide file tree
Showing 5 changed files with 255 additions and 50 deletions.
142 changes: 102 additions & 40 deletions crates/uv-resolver/src/resolution.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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;
Expand Down Expand Up @@ -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.
Expand All @@ -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.
Expand Down Expand Up @@ -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())?;

Expand All @@ -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(|| {
Expand All @@ -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())?;

Expand All @@ -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()),
Expand Down Expand Up @@ -259,6 +283,7 @@ impl ResolutionGraph {
Ok(Self {
petgraph,
hashes,
extras,
editables,
diagnostics,
})
Expand Down Expand Up @@ -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,
Expand All @@ -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(),
)
}
}

Expand All @@ -321,56 +355,78 @@ 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> {
Self {
resolution: underlying,
no_emit_packages,
show_hashes,
include_extras,
include_annotations,
annotation_style,
}
}
}

/// 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)]
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),
}

#[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,
}
}

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),
}
}
}

/// 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 {
// Collect all packages.
let mut nodes = self
.resolution
Expand All @@ -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))
})
Expand All @@ -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;
Expand Down
2 changes: 2 additions & 0 deletions crates/uv/src/commands/pip_compile.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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,
Expand Down Expand Up @@ -408,6 +409,7 @@ pub(crate) async fn pip_compile(
&resolution,
&no_emit_packages,
generate_hashes,
include_extras,
include_annotations,
annotation_style,
)
Expand Down
13 changes: 3 additions & 10 deletions crates/uv/src/compat/mod.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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>,
}
Expand Down Expand Up @@ -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() {
Expand Down
9 changes: 9 additions & 0 deletions crates/uv/src/main.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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,
Expand Down Expand Up @@ -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,
Expand Down
Loading

0 comments on commit ab99a18

Please sign in to comment.