Implement --emit-index-annotation to annotate source index for each package

crates/distribution-types/src/lib.rs

@@ -371,6 +371,13 @@ impl Dist {
         }
     }
 
+    pub fn index(&self) -> Option<String> {
+        match self {
+            Self::Built(dist) => dist.index(),
+            Self::Source(dist) => dist.index(),
+        }
+    }
+
     /// Returns the [`File`] instance, if this dist is from a registry with simple json api support
     pub fn file(&self) -> Option<&File> {
         match self {
@@ -388,6 +395,14 @@ impl Dist {
 }
 
 impl BuiltDist {
+    pub fn index(&self) -> Option<String> {
+        match self {
+            Self::Registry(registry) => Some(registry.index.to_string()),
+            Self::DirectUrl(_) => None,
+            Self::Path(_) => None,
+        }
+    }
+
     /// Returns the [`File`] instance, if this dist is from a registry with simple json api support
     pub fn file(&self) -> Option<&File> {
         match self {
@@ -406,6 +421,13 @@ impl BuiltDist {
 }
 
 impl SourceDist {
+    pub fn index(&self) -> Option<String> {
+        match self {
+            Self::Registry(registry) => Some(registry.index.to_string()),
+            Self::DirectUrl(_) | Self::Git(_) | Self::Path(_) => None,
+        }
+    }
+
     /// Returns the [`File`] instance, if this dist is from a registry with simple json api support
     pub fn file(&self) -> Option<&File> {
         match self {

crates/distribution-types/src/resolved.rs

@@ -31,6 +31,13 @@ impl ResolvedDist {
             Self::Installed(dist) => dist.is_editable(),
         }
     }
+
+    pub fn index(&self) -> Option<String> {
+        match self {
+            Self::Installable(dist) => dist.index(),
+            Self::Installed(_) => None,
+        }
+    }
 }
 
 impl ResolvedDistRef<'_> {

crates/uv-resolver/src/resolution.rs

@@ -493,6 +493,7 @@ impl ResolutionGraph {
 
 /// A [`std::fmt::Display`] implementation for the resolution graph.
 #[derive(Debug)]
+#[allow(clippy::struct_excessive_bools)]
 pub struct DisplayResolutionGraph<'a> {
     /// The underlying graph.
     resolution: &'a ResolutionGraph,
@@ -505,6 +506,8 @@ pub struct DisplayResolutionGraph<'a> {
     /// Whether to include annotations in the output, to indicate which dependency or dependencies
     /// requested each package.
     include_annotations: bool,
+    /// Whether to include indices in the output, to indicate which index was used for each package.
+    include_indices: bool,
     /// The style of annotation comments, used to indicate the dependencies that requested each
     /// package.
     annotation_style: AnnotationStyle,
@@ -518,19 +521,22 @@ impl<'a> From<&'a ResolutionGraph> for DisplayResolutionGraph<'a> {
             false,
             false,
             true,
+            false,
             AnnotationStyle::default(),
         )
     }
 }
 
 impl<'a> DisplayResolutionGraph<'a> {
     /// Create a new [`DisplayResolutionGraph`] for the given graph.
+    #[allow(clippy::fn_params_excessive_bools)]
     pub fn new(
         underlying: &'a ResolutionGraph,
         no_emit_packages: &'a [PackageName],
         show_hashes: bool,
         include_extras: bool,
         include_annotations: bool,
+        include_indices: bool,
         annotation_style: AnnotationStyle,
     ) -> DisplayResolutionGraph<'a> {
         Self {
@@ -539,6 +545,7 @@ impl<'a> DisplayResolutionGraph<'a> {
             show_hashes,
             include_extras,
             include_annotations,
+            include_indices,
             annotation_style,
         }
     }
@@ -576,6 +583,13 @@ impl<'a> Node<'a> {
             Node::Distribution(name, _, _) => NodeKey::Distribution(name),
         }
     }
+
+    fn index(&self) -> Option<String> {
+        match self {
+            Node::Editable(_, _) => None,
+            Node::Distribution(_, dist, _) => dist.index(),
+        }
+    }
 }
 
 impl Verbatim for Node<'_> {
@@ -714,6 +728,14 @@ impl std::fmt::Display for DisplayResolutionGraph<'_> {
                 // Write the line as is.
                 writeln!(f, "{line}")?;
             }
+
+            if self.include_indices && node.index().is_some() {
+                writeln!(
+                    f,
+                    "{}",
+                    format!("    # from {}", node.index().unwrap()).green()
+                )?;
+            }
         }
 
         Ok(())

crates/uv/src/commands/pip_compile.rs

@@ -63,6 +63,7 @@ pub(crate) async fn pip_compile(
     include_extras: bool,
     include_annotations: bool,
     include_header: bool,
+    include_indices: bool,
     custom_compile_command: Option<String>,
     include_index_url: bool,
     include_find_links: bool,
@@ -501,6 +502,7 @@ pub(crate) async fn pip_compile(
             generate_hashes,
             include_extras,
             include_annotations,
+            include_indices,
             annotation_style,
         )
     )?;

crates/uv/src/main.rs

@@ -335,6 +335,10 @@ struct PipCompileArgs {
     #[clap(long)]
     no_header: bool,
 
+    /// Include comment annotations indicating the index of each package.
+    #[clap(long)]
+    include_indices: bool,
+
     /// Change header comment to reflect custom command wrapping `uv pip compile`.
     #[clap(long, env = "UV_CUSTOM_COMPILE_COMMAND")]
     custom_compile_command: Option<String>,
@@ -1583,6 +1587,7 @@ async fn run() -> Result<ExitStatus> {
                 args.no_strip_extras,
                 !args.no_annotate,
                 !args.no_header,
+                args.include_indices,
                 args.custom_compile_command,
                 args.emit_index_url,
                 args.emit_find_links,

crates/uv/tests/pip_compile.rs

@@ -7334,3 +7334,73 @@ fn compile_index_url_fallback_prefer_primary() -> Result<()> {
 
     Ok(())
 }
+
+/// Ensure that `--include-indices` prints the index URL for each package.
+#[test]
+fn include_indices_pypi_org_simple() -> Result<()> {
+    let context = TestContext::new("3.12");
+
+    let requirements_in = context.temp_dir.child("requirements.in");
+    requirements_in.write_str("requests")?;
+
+    uv_snapshot!(context.compile()
+        .arg("requirements.in")
+        .arg("--include-indices"), @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 2024-03-25T00:00:00Z requirements.in --include-indices
+    certifi==2024.2.2
+        # via requests
+        # from https://pypi.org/simple
+    charset-normalizer==3.3.2
+        # via requests
+        # from https://pypi.org/simple
+    idna==3.6
+        # via requests
+        # from https://pypi.org/simple
+    requests==2.31.0
+        # from https://pypi.org/simple
+    urllib3==2.2.1
+        # via requests
+        # from https://pypi.org/simple
+
+    ----- stderr -----
+    Resolved 5 packages in [TIME]
+    "###
+    );
+
+    Ok(())
+}
+
+/// `--include-indices` where packages are pulled from two distinct indices.
+#[test]
+fn include_indices_multiple_indices() -> Result<()> {
+    let context = TestContext::new("3.12");
+
+    let requirements_in = context.temp_dir.child("requirements.in");
+    requirements_in.write_str("uv\nrequests")?;
+
+    uv_snapshot!(context.compile()
+        .arg("requirements.in")
+        .arg("--extra-index-url")
+        .arg("https://test.pypi.org/simple")
+        .arg("--include-indices"), @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 2024-03-25T00:00:00Z requirements.in --include-indices
+    requests==2.5.4.1
+        # from https://test.pypi.org/simple
+    uv==0.1.24
+        # from https://pypi.org/simple
+
+    ----- stderr -----
+    Resolved 2 packages in [TIME]
+    "###
+    );
+
+    Ok(())
+}