uv/crates/uv-resolver/src/lock/tree.rs

392 lines
14 KiB
Rust

use std::borrow::Cow;
use std::collections::BTreeSet;
use itertools::Itertools;
use rustc_hash::{FxHashMap, FxHashSet};
use uv_configuration::DevMode;
use uv_normalize::{ExtraName, GroupName, PackageName};
use uv_pypi_types::ResolverMarkerEnvironment;
use crate::lock::{Dependency, PackageId};
use crate::Lock;
#[derive(Debug)]
pub struct TreeDisplay<'env> {
/// The root nodes in the [`Lock`].
roots: Vec<&'env PackageId>,
/// The edges in the [`Lock`].
///
/// While the dependencies exist on the [`Lock`] directly, if `--invert` is enabled, the
/// direction must be inverted when constructing the tree.
dependencies: FxHashMap<&'env PackageId, Vec<Cow<'env, Dependency>>>,
optional_dependencies:
FxHashMap<&'env PackageId, FxHashMap<ExtraName, Vec<Cow<'env, Dependency>>>>,
dev_dependencies: FxHashMap<&'env PackageId, FxHashMap<GroupName, Vec<Cow<'env, Dependency>>>>,
/// Maximum display depth of the dependency tree.
depth: usize,
/// Whether to include development dependencies in the display.
dev: DevMode,
/// Prune the given packages from the display of the dependency tree.
prune: Vec<PackageName>,
/// Display only the specified packages.
packages: Vec<PackageName>,
/// Whether to de-duplicate the displayed dependencies.
no_dedupe: bool,
}
impl<'env> TreeDisplay<'env> {
/// Create a new [`DisplayDependencyGraph`] for the set of installed packages.
pub fn new(
lock: &'env Lock,
markers: Option<&'env ResolverMarkerEnvironment>,
depth: usize,
prune: Vec<PackageName>,
packages: Vec<PackageName>,
dev: DevMode,
no_dedupe: bool,
invert: bool,
) -> Self {
let mut non_roots = FxHashSet::default();
// Index all the dependencies. We could read these from the `Lock` directly, but we have to
// support `--invert`, so we might as well build them up in either case.
let mut dependencies: FxHashMap<_, Vec<_>> = FxHashMap::default();
let mut optional_dependencies: FxHashMap<_, FxHashMap<_, Vec<_>>> = FxHashMap::default();
let mut dev_dependencies: FxHashMap<_, FxHashMap<_, Vec<_>>> = FxHashMap::default();
for package in &lock.packages {
for dependency in &package.dependencies {
// Skip dependencies that don't apply to the current environment.
if let Some(environment_markers) = markers {
if !dependency
.complexified_marker
.evaluate(environment_markers, &[])
{
non_roots.insert(dependency.package_id.clone());
continue;
}
}
let parent = if invert {
&dependency.package_id
} else {
&package.id
};
let child = if invert {
Cow::Owned(Dependency {
package_id: package.id.clone(),
extra: dependency.extra.clone(),
simplified_marker: dependency.simplified_marker.clone(),
complexified_marker: dependency.complexified_marker.clone(),
})
} else {
Cow::Borrowed(dependency)
};
non_roots.insert(child.package_id.clone());
dependencies.entry(parent).or_default().push(child);
}
for (extra, dependencies) in &package.optional_dependencies {
for dependency in dependencies {
// Skip dependencies that don't apply to the current environment.
if let Some(environment_markers) = markers {
if !dependency
.complexified_marker
.evaluate(environment_markers, &[])
{
non_roots.insert(dependency.package_id.clone());
continue;
}
}
let parent = if invert {
&dependency.package_id
} else {
&package.id
};
let child = if invert {
Cow::Owned(Dependency {
package_id: package.id.clone(),
extra: dependency.extra.clone(),
simplified_marker: dependency.simplified_marker.clone(),
complexified_marker: dependency.complexified_marker.clone(),
})
} else {
Cow::Borrowed(dependency)
};
non_roots.insert(child.package_id.clone());
optional_dependencies
.entry(parent)
.or_default()
.entry(extra.clone())
.or_default()
.push(child);
}
}
for (group, dependencies) in &package.dev_dependencies {
for dependency in dependencies {
// Skip dependencies that don't apply to the current environment.
if let Some(environment_markers) = markers {
if !dependency
.complexified_marker
.evaluate(environment_markers, &[])
{
non_roots.insert(dependency.package_id.clone());
continue;
}
}
let parent = if invert {
&dependency.package_id
} else {
&package.id
};
let child = if invert {
Cow::Owned(Dependency {
package_id: package.id.clone(),
extra: dependency.extra.clone(),
simplified_marker: dependency.simplified_marker.clone(),
complexified_marker: dependency.complexified_marker.clone(),
})
} else {
Cow::Borrowed(dependency)
};
non_roots.insert(child.package_id.clone());
dev_dependencies
.entry(parent)
.or_default()
.entry(group.clone())
.or_default()
.push(child);
}
}
}
// Compute the root nodes.
let roots = lock
.packages
.iter()
.map(|dist| &dist.id)
.filter(|id| !non_roots.contains(*id))
.collect::<Vec<_>>();
Self {
roots,
dependencies,
optional_dependencies,
dev_dependencies,
depth,
dev,
prune,
packages,
no_dedupe,
}
}
/// Perform a depth-first traversal of the given package and its dependencies.
fn visit(
&'env self,
node: Node<'env>,
visited: &mut FxHashMap<&'env PackageId, Vec<&'env PackageId>>,
path: &mut Vec<&'env PackageId>,
) -> Vec<String> {
// Short-circuit if the current path is longer than the provided depth.
if path.len() > self.depth {
return Vec::new();
}
let line = {
let mut line = format!("{}", node.package_id().name);
if let Some(extras) = node.extras().filter(|extras| !extras.is_empty()) {
line.push_str(&format!("[{}]", extras.iter().join(",")));
}
line.push_str(&format!(" v{}", node.package_id().version));
match node {
Node::Root(_) => line,
Node::Dependency(_) => line,
Node::OptionalDependency(extra, _) => format!("{line} (extra: {extra})"),
Node::DevDependency(group, _) => format!("{line} (group: {group})"),
}
};
// Skip the traversal if:
// 1. The package is in the current traversal path (i.e., a dependency cycle).
// 2. The package has been visited and de-duplication is enabled (default).
if let Some(requirements) = visited.get(node.package_id()) {
if !self.no_dedupe || path.contains(&node.package_id()) {
return if requirements.is_empty() {
vec![line]
} else {
vec![format!("{} (*)", line)]
};
}
}
let dependencies: Vec<Node<'env>> = self
.dependencies
.get(node.package_id())
.filter(|_| self.dev != DevMode::Only)
.into_iter()
.flatten()
.map(|dep| Node::Dependency(dep.as_ref()))
.chain(
self.optional_dependencies
.get(node.package_id())
.filter(|_| self.dev != DevMode::Only)
.into_iter()
.flatten()
.flat_map(|(extra, deps)| {
deps.iter()
.map(move |dep| Node::OptionalDependency(extra, dep))
}),
)
.chain(
self.dev_dependencies
.get(node.package_id())
.filter(|_| self.dev != DevMode::Exclude)
.into_iter()
.flatten()
.flat_map(|(group, deps)| {
deps.iter().map(move |dep| Node::DevDependency(group, dep))
}),
)
.filter(|dep| !self.prune.contains(&dep.package_id().name))
.collect::<Vec<_>>();
let mut lines = vec![line];
// Keep track of the dependency path to avoid cycles.
visited.insert(
node.package_id(),
dependencies.iter().map(Node::package_id).collect(),
);
path.push(node.package_id());
for (index, dep) in dependencies.iter().enumerate() {
// For sub-visited packages, add the prefix to make the tree display user-friendly.
// The key observation here is you can group the tree as follows when you're at the
// root of the tree:
// root_package
// ├── level_1_0 // Group 1
// │ ├── level_2_0 ...
// │ │ ├── level_3_0 ...
// │ │ └── level_3_1 ...
// │ └── level_2_1 ...
// ├── level_1_1 // Group 2
// │ ├── level_2_2 ...
// │ └── level_2_3 ...
// └── level_1_2 // Group 3
// └── level_2_4 ...
//
// The lines in Group 1 and 2 have `├── ` at the top and `| ` at the rest while
// those in Group 3 have `└── ` at the top and ` ` at the rest.
// This observation is true recursively even when looking at the subtree rooted
// at `level_1_0`.
let (prefix_top, prefix_rest) = if dependencies.len() - 1 == index {
("└── ", " ")
} else {
("├── ", "")
};
for (visited_index, visited_line) in self.visit(*dep, visited, path).iter().enumerate()
{
let prefix = if visited_index == 0 {
prefix_top
} else {
prefix_rest
};
lines.push(format!("{prefix}{visited_line}"));
}
}
path.pop();
lines
}
/// Depth-first traverse the nodes to render the tree.
fn render(&self) -> Vec<String> {
let mut visited = FxHashMap::default();
let mut path = Vec::new();
let mut lines = Vec::new();
if self.packages.is_empty() {
for id in &self.roots {
path.clear();
lines.extend(self.visit(Node::Root(id), &mut visited, &mut path));
}
} else {
let by_package: FxHashMap<_, _> = self.roots.iter().map(|id| (&id.name, id)).collect();
for package in &self.packages {
if let Some(id) = by_package.get(package) {
path.clear();
lines.extend(self.visit(Node::Root(id), &mut visited, &mut path));
}
}
}
lines
}
}
#[derive(Debug, Copy, Clone)]
enum Node<'env> {
Root(&'env PackageId),
Dependency(&'env Dependency),
OptionalDependency(&'env ExtraName, &'env Dependency),
DevDependency(&'env GroupName, &'env Dependency),
}
impl<'env> Node<'env> {
fn package_id(&self) -> &'env PackageId {
match self {
Self::Root(id) => id,
Self::Dependency(dep) => &dep.package_id,
Self::OptionalDependency(_, dep) => &dep.package_id,
Self::DevDependency(_, dep) => &dep.package_id,
}
}
fn extras(&self) -> Option<&BTreeSet<ExtraName>> {
match self {
Self::Root(_) => None,
Self::Dependency(dep) => Some(&dep.extra),
Self::OptionalDependency(_, dep) => Some(&dep.extra),
Self::DevDependency(_, dep) => Some(&dep.extra),
}
}
}
impl std::fmt::Display for TreeDisplay<'_> {
fn fmt(&self, f: &mut std::fmt::Formatter) -> std::fmt::Result {
use owo_colors::OwoColorize;
let mut deduped = false;
for line in self.render() {
deduped |= line.contains('*');
writeln!(f, "{line}")?;
}
if deduped {
let message = if self.no_dedupe {
"(*) Package tree is a cycle and cannot be shown".italic()
} else {
"(*) Package tree already displayed".italic()
};
writeln!(f, "{message}")?;
}
Ok(())
}
}