From 7cd98d24995122ffe58c41f05e7de28cf9747d60 Mon Sep 17 00:00:00 2001 From: Chan Kang Date: Wed, 10 Apr 2024 12:05:58 -0400 Subject: [PATCH] Implement `--emit-index-annotation` to annotate source index for each package (#2926) ## Summary resolves https://github.com/astral-sh/uv/issues/2852 ## Test Plan add a couple of tests: - one covering the simplest case with all packages pulled from a single index. - another where packages are pull from two distinct indices. tested manually as well: ``` $ (echo 'pandas'; echo 'torch') | UV_EXTRA_INDEX_URL='https://download.pytorch.org/whl/cpu' cargo run pip compile - --include-indices Finished dev [unoptimized + debuginfo] target(s) in 0.60s Running `target/debug/uv pip compile - --include-indices` Resolved 15 packages in 686ms # This file was autogenerated by uv via the following command: # uv pip compile - --include-indices filelock==3.9.0 # via torch # from https://download.pytorch.org/whl/cpu fsspec==2023.4.0 # via torch # from https://download.pytorch.org/whl/cpu jinja2==3.1.2 # via torch # from https://download.pytorch.org/whl/cpu markupsafe==2.1.3 # via jinja2 # from https://download.pytorch.org/whl/cpu mpmath==1.3.0 # via sympy # from https://download.pytorch.org/whl/cpu networkx==3.2.1 # via torch # from https://download.pytorch.org/whl/cpu numpy==1.26.3 # via pandas # from https://download.pytorch.org/whl/cpu pandas==2.2.1 # from https://pypi.org/simple python-dateutil==2.9.0.post0 # via pandas # from https://pypi.org/simple pytz==2024.1 # via pandas # from https://pypi.org/simple six==1.16.0 # via python-dateutil # from https://pypi.org/simple sympy==1.12 # via torch # from https://download.pytorch.org/whl/cpu torch==2.2.2 # from https://download.pytorch.org/whl/cpu typing-extensions==4.8.0 # via torch # from https://download.pytorch.org/whl/cpu tzdata==2024.1 # via pandas # from https://pypi.org/simple ``` --- crates/distribution-types/src/lib.rs | 27 +++- crates/distribution-types/src/resolved.rs | 12 +- crates/uv-resolver/src/resolution.rs | 27 +++- crates/uv/src/commands/pip_compile.rs | 2 + crates/uv/src/main.rs | 12 +- crates/uv/tests/pip_compile.rs | 146 ++++++++++++++++++++++ 6 files changed, 219 insertions(+), 7 deletions(-) diff --git a/crates/distribution-types/src/lib.rs b/crates/distribution-types/src/lib.rs index 4b84e6e38..6c2ad6c81 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 91015e4f3..1f72ba4a0 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 fe81b49b3..9dbf758dc 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; @@ -499,6 +499,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, @@ -511,6 +512,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, @@ -524,6 +527,7 @@ impl<'a> From<&'a ResolutionGraph> for DisplayResolutionGraph<'a> { false, false, true, + false, AnnotationStyle::default(), ) } @@ -531,12 +535,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 { @@ -545,6 +551,7 @@ impl<'a> DisplayResolutionGraph<'a> { show_hashes, include_extras, include_annotations, + include_index_annotation, annotation_style, } } @@ -582,6 +589,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<'_> { @@ -666,6 +681,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 @@ -720,6 +737,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 a635c6ed0..b48d33e60 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 22f318c7f..5b6ec7ebf 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 79281c57a..2c153868e 100644 --- a/crates/uv/tests/pip_compile.rs +++ b/crates/uv/tests/pip_compile.rs @@ -7430,3 +7430,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(()) +}