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
```
This commit is contained in:
Chan Kang 2024-04-10 12:05:58 -04:00 committed by GitHub
parent a01143980a
commit 7cd98d2499
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
6 changed files with 219 additions and 7 deletions

View file

@ -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 {

View file

@ -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<'_> {

View file

@ -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(())

View file

@ -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,
)
)?;

View file

@ -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,

View file

@ -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(())
}