Repository: astral-sh/uv (Rust)

Issue:
Annotate the index from which a package is pulled in `uv pip compile`
Awaiting that #171 gets resolved, it would be very useful if `uv pip compile` could generate annotations with the “source” index.

Currently `uv pip compile --extra-index-url <MYINDEX> --no-header <(echo MYPACKAGE)` outputs something like
```
DEPENDENCY==1.18.5
    # via MYPACKAGE
MYPACKAGE==0.0.5
```
and I found no easy way (apart from analyzing the `--verbose` trace or cross-referencing the hashes from `--generate-hashes`) to understand which package comes from which index.

I would like explicit annotations, e.g.
```
DEPENDENCY==1.18.5
    # via MYPACKAGE
    # from https://pypi.org/simple
MYPACKAGE==0.0.5
    # from <MYINDEX>
```



Discussion: Technically, we could do this without #171. We know which index each package comes from at compile-time.
> Technically, we could do this without #171. We know which index each package comes from at compile-time.

Sorry for my unclear wording: indeed I'm suggesting to implent this enhancement as a mitigation for possible dependency confusion attacks *before* #171 is finalized.

Knowing the origin of a package without the need to trace the index-strategy could also be helpful for users that are not fully aware of the differences with `pip`. 
Makes sense :)
Ah cool makes sense. As long as there's not an expectation that we will guarantee the given index is used on install.




Fix (patch):
diff --git a/crates/distribution-types/src/lib.rs b/crates/distribution-types/src/lib.rs
index 4b84e6e38a704..6c2ad6c81c458 100644
--- a/crates/distribution-types/src/lib.rs
+++ b/crates/distribution-types/src/lib.rs
@@ -371,6 +371,14 @@ impl Dist {
         }
     }
 
+    /// Returns the [`IndexUrl`], if the distribution is from a registry.
+    pub fn index(&self) -> Option<&IndexUrl> {
+        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,7 +396,16 @@ impl Dist {
 }
 
 impl BuiltDist {
-    /// Returns the [`File`] instance, if this dist is from a registry with simple json api support
+    /// Returns the [`IndexUrl`], if the distribution is from a registry.
+    pub fn index(&self) -> Option<&IndexUrl> {
+        match self {
+            Self::Registry(registry) => Some(&registry.index),
+            Self::DirectUrl(_) => None,
+            Self::Path(_) => None,
+        }
+    }
+
+    /// Returns the [`File`] instance, if this distribution is from a registry.
     pub fn file(&self) -> Option<&File> {
         match self {
             Self::Registry(registry) => Some(&registry.file),
@@ -406,6 +423,14 @@ impl BuiltDist {
 }
 
 impl SourceDist {
+    /// Returns the [`IndexUrl`], if the distribution is from a registry.
+    pub fn index(&self) -> Option<&IndexUrl> {
+        match self {
+            Self::Registry(registry) => Some(&registry.index),
+            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 {
diff --git a/crates/distribution-types/src/resolved.rs b/crates/distribution-types/src/resolved.rs
index 91015e4f389da..1f72ba4a08def 100644
--- a/crates/distribution-types/src/resolved.rs
+++ b/crates/distribution-types/src/resolved.rs
@@ -3,8 +3,8 @@ use std::fmt::{Display, Formatter};
 use pep508_rs::PackageName;
 
 use crate::{
-    Dist, DistributionId, DistributionMetadata, Identifier, InstalledDist, Name, ResourceId,
-    VersionOrUrl,
+    Dist, DistributionId, DistributionMetadata, Identifier, IndexUrl, InstalledDist, Name,
+    ResourceId, VersionOrUrl,
 };
 
 /// A distribution that can be used for resolution and installation.
@@ -31,6 +31,14 @@ impl ResolvedDist {
             Self::Installed(dist) => dist.is_editable(),
         }
     }
+
+    /// Returns the [`IndexUrl`], if the distribution is from a registry.
+    pub fn index(&self) -> Option<&IndexUrl> {
+        match self {
+            Self::Installable(dist) => dist.index(),
+            Self::Installed(_) => None,
+        }
+    }
 }
 
 impl ResolvedDistRef<'_> {
diff --git a/crates/uv-resolver/src/resolution.rs b/crates/uv-resolver/src/resolution.rs
index da93cce163346..533fcd0dad7fc 100644
--- a/crates/uv-resolver/src/resolution.rs
+++ b/crates/uv-resolver/src/resolution.rs
@@ -12,7 +12,7 @@ use pubgrub::type_aliases::SelectedDependencies;
 use rustc_hash::{FxHashMap, FxHashSet};
 
 use distribution_types::{
-    Dist, DistributionMetadata, LocalEditable, Name, PackageId, ResolvedDist, Verbatim,
+    Dist, DistributionMetadata, IndexUrl, LocalEditable, Name, PackageId, ResolvedDist, Verbatim,
     VersionOrUrl,
 };
 use once_map::OnceMap;
@@ -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 indexes in the output, to indicate which index was used for each package.
+    include_index_annotation: bool,
     /// The style of annotation comments, used to indicate the dependencies that requested each
     /// package.
     annotation_style: AnnotationStyle,
@@ -518,6 +521,7 @@ impl<'a> From<&'a ResolutionGraph> for DisplayResolutionGraph<'a> {
             false,
             false,
             true,
+            false,
             AnnotationStyle::default(),
         )
     }
@@ -525,12 +529,14 @@ impl<'a> From<&'a ResolutionGraph> for DisplayResolutionGraph<'a> {
 
 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_index_annotation: bool,
         annotation_style: AnnotationStyle,
     ) -> DisplayResolutionGraph<'a> {
         Self {
@@ -539,6 +545,7 @@ impl<'a> DisplayResolutionGraph<'a> {
             show_hashes,
             include_extras,
             include_annotations,
+            include_index_annotation,
             annotation_style,
         }
     }
@@ -576,6 +583,14 @@ impl<'a> Node<'a> {
             Node::Distribution(name, _, _) => NodeKey::Distribution(name),
         }
     }
+
+    /// Return the [`IndexUrl`] of the distribution, if any.
+    fn index(&self) -> Option<&IndexUrl> {
+        match self {
+            Node::Editable(_, _) => None,
+            Node::Distribution(_, dist, _) => dist.index(),
+        }
+    }
 }
 
 impl Verbatim for Node<'_> {
@@ -660,6 +675,8 @@ impl std::fmt::Display for DisplayResolutionGraph<'_> {
             // Determine the annotation comment and separator (between comment and requirement).
             let mut annotation = None;
 
+            // If enabled, include annotations to indicate the dependencies that requested each
+            // package (e.g., `# via mypy`).
             if self.include_annotations {
                 // Display all dependencies.
                 let mut edges = self
@@ -714,6 +731,14 @@ impl std::fmt::Display for DisplayResolutionGraph<'_> {
                 // Write the line as is.
                 writeln!(f, "{line}")?;
             }
+
+            // If enabled, include indexes to indicate which index was used for each package (e.g.,
+            // `# from https://pypi.org/simple`).
+            if self.include_index_annotation {
+                if let Some(index) = node.index() {
+                    writeln!(f, "{}", format!("    # from {index}").green())?;
+                }
+            }
         }
 
         Ok(())
diff --git a/crates/uv/src/commands/pip_compile.rs b/crates/uv/src/commands/pip_compile.rs
index a635c6ed008a1..b48d33e60e185 100644
--- a/crates/uv/src/commands/pip_compile.rs
+++ b/crates/uv/src/commands/pip_compile.rs
@@ -67,6 +67,7 @@ pub(crate) async fn pip_compile(
     include_index_url: bool,
     include_find_links: bool,
     include_marker_expression: bool,
+    include_index_annotation: bool,
     index_locations: IndexLocations,
     index_strategy: IndexStrategy,
     keyring_provider: KeyringProvider,
@@ -501,6 +502,7 @@ pub(crate) async fn pip_compile(
             generate_hashes,
             include_extras,
             include_annotations,
+            include_index_annotation,
             annotation_style,
         )
     )?;
diff --git a/crates/uv/src/main.rs b/crates/uv/src/main.rs
index 134e11a30f6c1..8f30598ca54c5 100644
--- a/crates/uv/src/main.rs
+++ b/crates/uv/src/main.rs
@@ -335,6 +335,10 @@ struct PipCompileArgs {
     #[clap(long)]
     no_header: bool,
 
+    /// Choose the style of the annotation comments, which indicate the source of each package.
+    #[clap(long, default_value_t=AnnotationStyle::Split, value_enum)]
+    annotation_style: AnnotationStyle,
+
     /// Change header comment to reflect custom command wrapping `uv pip compile`.
     #[clap(long, env = "UV_CUSTOM_COMPILE_COMMAND")]
     custom_compile_command: Option<String>,
@@ -495,9 +499,10 @@ struct PipCompileArgs {
     #[clap(long, hide = true)]
     emit_marker_expression: bool,
 
-    /// Choose the style of the annotation comments, which indicate the source of each package.
-    #[clap(long, default_value_t=AnnotationStyle::Split, value_enum)]
-    annotation_style: AnnotationStyle,
+    /// Include comment annotations indicating the index used to resolve each package (e.g.,
+    /// `# from https://pypi.org/simple`).
+    #[clap(long)]
+    emit_index_annotation: bool,
 
     #[command(flatten)]
     compat_args: compat::PipCompileCompatArgs,
@@ -1587,6 +1592,7 @@ async fn run() -> Result<ExitStatus> {
                 args.emit_index_url,
                 args.emit_find_links,
                 args.emit_marker_expression,
+                args.emit_index_annotation,
                 index_urls,
                 args.index_strategy,
                 args.keyring_provider,