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

@@ -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,37 +363,35 @@ impl<'a> DisplayResolutionGraph<'a> {
             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.
+#[derive(Debug)]
-impl std::fmt::Display for DisplayResolutionGraph<'_> {
+enum Node<'a> {
-    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),
+    Distribution(&'a PackageName, &'a Dist, &'a [ExtraName]),
-        }
+}
 
-        #[derive(Debug, PartialEq, Eq, PartialOrd, Ord)]
+#[derive(Debug, PartialEq, Eq, PartialOrd, Ord)]
-        enum NodeKey<'a> {
+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> {
+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,
+            Node::Distribution(name, _, _) => name,
         }
     }
 
@@ -366,11 +399,34 @@ impl std::fmt::Display for DisplayResolutionGraph<'_> {
     fn key(&self) -> NodeKey<'a> {
         match self {
             Node::Editable(_, editable) => NodeKey::Editable(editable.verbatim()),
-                    Node::Distribution(name, _) => NodeKey::Distribution(name),
+            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
@@ -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 {
+            let mut line = node.verbatim().to_string();
-                Node::Distribution(_, dist) => format!("{}", dist.verbatim()),
-                Node::Editable(_, editable) => format!("-e {}", editable.verbatim()),
-            };
 
             // Display the distribution hashes, if any.
             let mut has_hashes = false;

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

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<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).");
+            warn_user!(
-        }
+                "pip-compile's `--strip-extras` has no effect (uv strips extras by default)."
-
+            );
-        if self.no_strip_extras {
-            return Err(anyhow!(
-                "pip-compile's `--no-strip-extras` is unsupported (uv always strips extras)."
-            ));
         }
 
         if self.pip_args.is_some() {

crates/uv/src/main.rs

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

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]