From 61014d48b08b9f727255efe5d3db5612b8e613b3 Mon Sep 17 00:00:00 2001 From: Chan Kang Date: Mon, 1 Jul 2024 17:12:59 -0400 Subject: [PATCH] Implement `--package` for `pip tree` (#4655) ## Summary Part of https://github.com/astral-sh/uv/issues/4439. ## Test Plan The existing tests pass + added a couple of tests to ensure `--package` behaves as expected. --- crates/uv-cli/src/lib.rs | 5 + crates/uv/src/commands/pip/tree.rs | 35 +++++-- crates/uv/src/main.rs | 1 + crates/uv/src/settings.rs | 3 + crates/uv/tests/pip_tree.rs | 150 +++++++++++++++++++++++++++++ 5 files changed, 185 insertions(+), 9 deletions(-) diff --git a/crates/uv-cli/src/lib.rs b/crates/uv-cli/src/lib.rs index dbc30b54a068..31b4fcf52b1f 100644 --- a/crates/uv-cli/src/lib.rs +++ b/crates/uv-cli/src/lib.rs @@ -1428,6 +1428,11 @@ pub struct PipTreeArgs { /// Prune the given package from the display of the dependency tree. #[arg(long)] pub prune: Vec, + + /// Display only the specified packages. + #[arg(long)] + pub package: Vec, + /// Do not de-duplicate repeated dependencies. /// Usually, when a package has already displayed its dependencies, /// further occurrences will not re-display its dependencies, diff --git a/crates/uv/src/commands/pip/tree.rs b/crates/uv/src/commands/pip/tree.rs index f1e70cdb7eb3..888fa7e7fbc2 100644 --- a/crates/uv/src/commands/pip/tree.rs +++ b/crates/uv/src/commands/pip/tree.rs @@ -25,6 +25,7 @@ use crate::printer::Printer; pub(crate) fn pip_tree( depth: u8, prune: Vec, + package: Vec, no_dedupe: bool, invert: bool, strict: bool, @@ -53,6 +54,7 @@ pub(crate) fn pip_tree( &site_packages, depth.into(), prune, + package, no_dedupe, invert, environment.interpreter().markers(), @@ -117,6 +119,8 @@ struct DisplayDependencyGraph<'env> { depth: usize, /// Prune the given packages from the display of the dependency tree. prune: Vec, + /// Display only the specified packages. + package: Vec, /// Whether to de-duplicate the displayed dependencies. no_dedupe: bool, /// Map from package name to its requirements. @@ -131,6 +135,7 @@ impl<'env> DisplayDependencyGraph<'env> { site_packages: &'env SitePackages, depth: usize, prune: Vec, + package: Vec, no_dedupe: bool, invert: bool, markers: &'env MarkerEnvironment, @@ -158,6 +163,7 @@ impl<'env> DisplayDependencyGraph<'env> { site_packages, depth, prune, + package, no_dedupe, requirements, }) @@ -260,17 +266,28 @@ impl<'env> DisplayDependencyGraph<'env> { let mut path: Vec<&PackageName> = Vec::new(); let mut lines: Vec = Vec::new(); - // The root nodes are those that are not required by any other package. - let children: HashSet<_> = self.requirements.values().flatten().collect(); - for site_package in self.site_packages.iter() { - // If the current package is not required by any other package, start the traversal - // with the current package as the root. - if !children.contains(site_package.name()) { - path.clear(); - lines.extend(self.visit(site_package, &mut visited, &mut path)?); + if self.package.is_empty() { + // The root nodes are those that are not required by any other package. + let children: HashSet<_> = self.requirements.values().flatten().collect(); + for site_package in self.site_packages.iter() { + // If the current package is not required by any other package, start the traversal + // with the current package as the root. + if !children.contains(site_package.name()) { + path.clear(); + lines.extend(self.visit(site_package, &mut visited, &mut path)?); + } + } + } else { + for (index, package) in self.package.iter().enumerate() { + if index != 0 { + lines.push(String::new()); + } + for installed_dist in self.site_packages.get_packages(package) { + path.clear(); + lines.extend(self.visit(installed_dist, &mut visited, &mut path)?); + } } } - Ok(lines) } } diff --git a/crates/uv/src/main.rs b/crates/uv/src/main.rs index 5fc76256f7d5..f265057bfc88 100644 --- a/crates/uv/src/main.rs +++ b/crates/uv/src/main.rs @@ -550,6 +550,7 @@ async fn run() -> Result { commands::pip_tree( args.depth, args.prune, + args.package, args.no_dedupe, args.invert, args.shared.strict, diff --git a/crates/uv/src/settings.rs b/crates/uv/src/settings.rs index aa957ce48c54..7b907c229cbd 100644 --- a/crates/uv/src/settings.rs +++ b/crates/uv/src/settings.rs @@ -1066,6 +1066,7 @@ impl PipShowSettings { pub(crate) struct PipTreeSettings { pub(crate) depth: u8, pub(crate) prune: Vec, + pub(crate) package: Vec, pub(crate) no_dedupe: bool, pub(crate) invert: bool, // CLI-only settings. @@ -1078,6 +1079,7 @@ impl PipTreeSettings { let PipTreeArgs { depth, prune, + package, no_dedupe, invert, strict, @@ -1091,6 +1093,7 @@ impl PipTreeSettings { Self { depth, prune, + package, no_dedupe, invert, // Shared settings. diff --git a/crates/uv/tests/pip_tree.rs b/crates/uv/tests/pip_tree.rs index c490b32341ea..e3eca62e130f 100644 --- a/crates/uv/tests/pip_tree.rs +++ b/crates/uv/tests/pip_tree.rs @@ -1354,3 +1354,153 @@ fn with_editable() { "### ); } + +#[test] +#[cfg(target_os = "macos")] +fn package_flag_complex() { + let context = TestContext::new("3.12"); + + let requirements_txt = context.temp_dir.child("requirements.txt"); + requirements_txt.write_str("packse").unwrap(); + + uv_snapshot!(context + .pip_install() + .arg("-r") + .arg("requirements.txt") + .arg("--strict"), @r###" + success: true + exit_code: 0 + ----- stdout ----- + + ----- stderr ----- + Resolved 32 packages in [TIME] + Prepared 32 packages in [TIME] + Installed 32 packages in [TIME] + + certifi==2024.2.2 + + charset-normalizer==3.3.2 + + chevron-blue==0.2.1 + + docutils==0.20.1 + + hatchling==1.22.4 + + idna==3.6 + + importlib-metadata==7.1.0 + + jaraco-classes==3.3.1 + + jaraco-context==4.3.0 + + jaraco-functools==4.0.0 + + keyring==25.0.0 + + markdown-it-py==3.0.0 + + mdurl==0.1.2 + + more-itertools==10.2.0 + + msgspec==0.18.6 + + nh3==0.2.15 + + packaging==24.0 + + packse==0.3.12 + + pathspec==0.12.1 + + pkginfo==1.10.0 + + pluggy==1.4.0 + + pygments==2.17.2 + + readme-renderer==43.0 + + requests==2.31.0 + + requests-toolbelt==1.0.0 + + rfc3986==2.0.0 + + rich==13.7.1 + + setuptools==69.2.0 + + trove-classifiers==2024.3.3 + + twine==4.0.2 + + urllib3==2.2.1 + + zipp==3.18.1 + "### + ); + + uv_snapshot!( + context.filters(), + context.pip_tree() + .arg("--package") + .arg("hatchling") + .arg("--package") + .arg("keyring"), @r###" + success: true + exit_code: 0 + ----- stdout ----- + hatchling v1.22.4 + ├── packaging v24.0 + ├── pathspec v0.12.1 + ├── pluggy v1.4.0 + └── trove-classifiers v2024.3.3 + + keyring v25.0.0 + ├── jaraco-classes v3.3.1 + │ └── more-itertools v10.2.0 + ├── jaraco-functools v4.0.0 + │ └── more-itertools v10.2.0 + └── jaraco-context v4.3.0 + + ----- stderr ----- + "### + ); +} + +#[test] +fn package_flag() { + let context = TestContext::new("3.12"); + + let requirements_txt = context.temp_dir.child("requirements.txt"); + requirements_txt + .write_str("scikit-learn==1.4.1.post1") + .unwrap(); + + uv_snapshot!(context + .pip_install() + .arg("-r") + .arg("requirements.txt") + .arg("--strict"), @r###" + success: true + exit_code: 0 + ----- stdout ----- + + ----- stderr ----- + Resolved 5 packages in [TIME] + Prepared 5 packages in [TIME] + Installed 5 packages in [TIME] + + joblib==1.3.2 + + numpy==1.26.4 + + scikit-learn==1.4.1.post1 + + scipy==1.12.0 + + threadpoolctl==3.4.0 + "### + ); + + uv_snapshot!( + context.filters(), + context.pip_tree() + .arg("--package") + .arg("numpy"), + @r###" + success: true + exit_code: 0 + ----- stdout ----- + numpy v1.26.4 + + ----- stderr ----- + "### + ); + + uv_snapshot!( + context.filters(), + context.pip_tree() + .arg("--package") + .arg("scipy") + .arg("--package") + .arg("joblib"), + @r###" + success: true + exit_code: 0 + ----- stdout ----- + scipy v1.12.0 + └── numpy v1.26.4 + + joblib v1.3.2 + + ----- stderr ----- + "### + ); +}