add index in pip compile output

astral-sh · Apr 9, 2024 · 280f4bc · 280f4bc
1 parent 06e96a8
commit 280f4bc
Show file tree

Hide file tree

Showing 3 changed files with 48 additions and 3 deletions.
diff --git a/crates/distribution-types/src/lib.rs b/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 {

diff --git a/crates/distribution-types/src/resolved.rs b/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<'_> {

diff --git a/crates/uv-resolver/src/resolution.rs b/crates/uv-resolver/src/resolution.rs
@@ -576,6 +576,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<'_> {
@@ -672,6 +679,12 @@ impl std::fmt::Display for DisplayResolutionGraph<'_> {
                     .collect::<Vec<_>>();
                 edges.sort_unstable_by_key(|package| package.name());
 
+                let index = node.index();
+                let index_line = if index.is_some() {
+                    format!("\n    # from {}", index.unwrap())
+                } else {
+                    "".to_string()
+                };
                 match self.annotation_style {
                     AnnotationStyle::Line => {
                         if !edges.is_empty() {
@@ -681,15 +694,17 @@ impl std::fmt::Display for DisplayResolutionGraph<'_> {
                                 .map(|dependency| dependency.name().to_string())
                                 .collect::<Vec<_>>()
                                 .join(", ");
-                            let comment = format!("# via {deps}").green().to_string();
+                            let comment = format!("# via {deps}{index_line}").green().to_string();
                             annotation = Some((separator, comment));
                         }
                     }
                     AnnotationStyle::Split => match edges.as_slice() {
                         [] => {}
                         [edge] => {
                             let separator = "\n";
-                            let comment = format!("    # via {}", edge.name()).green().to_string();
+                            let comment = format!("    # via {}{index_line}", edge.name())
+                                .green()
+                                .to_string();
                             annotation = Some((separator, comment));
                         }
                         edges => {
@@ -699,7 +714,8 @@ impl std::fmt::Display for DisplayResolutionGraph<'_> {
                                 .map(|dependency| format!("    #   {}", dependency.name()))
                                 .collect::<Vec<_>>()
                                 .join("\n");
-                            let comment = format!("    # via\n{deps}").green().to_string();
+                            let comment =
+                                format!("    # via\n{deps}{index_line}").green().to_string();
                             annotation = Some((separator, comment));
                         }
                     },