diff --git a/crates/distribution-types/src/lib.rs b/crates/distribution-types/src/lib.rs index 4b84e6e38a70..6c2ad6c81c45 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(®istry.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(®istry.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(®istry.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 91015e4f389d..1f72ba4a08de 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 da93cce16334..533fcd0dad7f 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 a635c6ed008a..b48d33e60e18 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 134e11a30f6c..8f30598ca54c 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, @@ -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 { args.emit_index_url, args.emit_find_links, args.emit_marker_expression, + args.emit_index_annotation, index_urls, args.index_strategy, args.keyring_provider, diff --git a/crates/uv/tests/pip_compile.rs b/crates/uv/tests/pip_compile.rs index 292d4507bdfe..1f82bc188738 100644 --- a/crates/uv/tests/pip_compile.rs +++ b/crates/uv/tests/pip_compile.rs @@ -7334,3 +7334,149 @@ fn compile_index_url_fallback_prefer_primary() -> Result<()> { Ok(()) } + +/// Ensure that `--emit-index-annotation` prints the index URL for each package. +#[test] +fn emit_index_annotation_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("--emit-index-annotation"), @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 --emit-index-annotation + 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(()) +} + +/// Ensure that `--emit-index-annotation` plays nicely with `--no-annotate`. +/// +/// For now, `--no-annotate` doesn't affect `--emit-index-annotation`, in that we still emit the +/// index annotation, and leave `--no-annotate` to only affect the package _source_ annotations. +#[test] +fn emit_index_annotation_no_annotate() -> 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("--emit-index-annotation") + .arg("--no-annotate"), @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 --emit-index-annotation --no-annotate + certifi==2024.2.2 + # from https://pypi.org/simple + charset-normalizer==3.3.2 + # from https://pypi.org/simple + idna==3.6 + # from https://pypi.org/simple + requests==2.31.0 + # from https://pypi.org/simple + urllib3==2.2.1 + # from https://pypi.org/simple + + ----- stderr ----- + Resolved 5 packages in [TIME] + "### + ); + + Ok(()) +} + +/// Ensure that `--emit-index-annotation` plays nicely with `--annotation-style=line`. +#[test] +fn emit_index_annotation_line() -> 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("--emit-index-annotation") + .arg("--annotation-style") + .arg("line"), @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 --emit-index-annotation --annotation-style line + 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(()) +} + +/// `--emit-index-annotation` where packages are pulled from two distinct indexes. +#[test] +fn emit_index_annotation_multiple_indexes() -> 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("--emit-index-annotation"), @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 --emit-index-annotation + 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(()) +}