1use std::cmp::Ordering;
2use std::collections::{BTreeSet, VecDeque};
3use std::fmt::Write;
4
5use either::Either;
6use itertools::Itertools;
7use owo_colors::OwoColorize;
8use petgraph::graph::{EdgeIndex, NodeIndex};
9use petgraph::prelude::EdgeRef;
10use petgraph::{Direction, Graph};
11use rustc_hash::{FxBuildHasher, FxHashMap, FxHashSet};
12
13use uv_configuration::DependencyGroupsWithDefaults;
14use uv_console::human_readable_bytes;
15use uv_normalize::{ExtraName, GroupName, PackageName};
16use uv_pep440::Version;
17use uv_pep508::MarkerTree;
18use uv_pypi_types::ResolverMarkerEnvironment;
19
20use crate::lock::{Package, PackageId};
21use crate::{ConflictMarker, Lock, PackageMap, UniversalMarker};
22
23#[derive(Debug)]
24pub struct TreeDisplay<'env> {
25 graph: petgraph::graph::Graph<Node<'env>, Edge<'env>, petgraph::Directed>,
27 roots: Vec<NodeIndex>,
29 latest: &'env PackageMap<Version>,
31 depth: usize,
33 no_dedupe: bool,
35 invert: bool,
37 lock: &'env Lock,
39 show_sizes: bool,
41 conflict_marker: UniversalMarker,
43}
44
45impl<'env> TreeDisplay<'env> {
46 pub fn new(
48 lock: &'env Lock,
49 markers: Option<&'env ResolverMarkerEnvironment>,
50 latest: &'env PackageMap<Version>,
51 depth: usize,
52 prune: &[PackageName],
53 packages: &[PackageName],
54 groups: &DependencyGroupsWithDefaults,
55 no_dedupe: bool,
56 invert: bool,
57 show_sizes: bool,
58 ) -> Self {
59 let members: BTreeSet<&PackageId> = if lock.members().is_empty() {
67 lock.root().into_iter().map(|package| &package.id).collect()
68 } else {
69 lock.packages
70 .iter()
71 .filter_map(|package| {
72 if lock.members().contains(&package.id.name) {
73 Some(&package.id)
74 } else {
75 None
76 }
77 })
78 .collect()
79 };
80
81 let conflict_marker = UniversalMarker::new(
84 MarkerTree::TRUE,
85 ConflictMarker::from_conflicts(lock.conflicts()),
86 );
87
88 let size_guess = lock.packages.len();
90 let mut graph =
91 Graph::<Node, Edge, petgraph::Directed>::with_capacity(size_guess, size_guess);
92 let mut inverse = FxHashMap::with_capacity_and_hasher(size_guess, FxBuildHasher);
93 let mut queue: VecDeque<(&PackageId, Option<&ExtraName>)> = VecDeque::new();
94 let mut seen = FxHashSet::default();
95
96 let root = graph.add_node(Node::Root);
97
98 for id in members.iter().copied() {
100 if prune.contains(&id.name) {
101 continue;
102 }
103
104 let dist = lock.find_by_id(id);
105
106 let index = *inverse
110 .entry(id)
111 .or_insert_with(|| graph.add_node(Node::Package(id)));
112
113 graph.add_edge(root, index, Edge::Prod(None, UniversalMarker::TRUE));
115
116 if groups.prod() {
117 if seen.insert((id, None)) {
119 queue.push_back((id, None));
120 }
121
122 for extra in dist.optional_dependencies.keys() {
124 if seen.insert((id, Some(extra))) {
125 queue.push_back((id, Some(extra)));
126 }
127 }
128 }
129
130 for (group, dep) in dist
132 .dependency_groups
133 .iter()
134 .filter_map(|(group, deps)| {
135 if groups.contains(group) {
136 Some(deps.iter().map(move |dep| (group, dep)))
137 } else {
138 None
139 }
140 })
141 .flatten()
142 {
143 if prune.contains(&dep.package_id.name) {
144 continue;
145 }
146
147 if markers
148 .is_some_and(|markers| !dep.complexified_marker.evaluate_no_extras(markers))
149 {
150 continue;
151 }
152
153 let dep_index = *inverse
155 .entry(&dep.package_id)
156 .or_insert_with(|| graph.add_node(Node::Package(&dep.package_id)));
157
158 graph.add_edge(
160 index,
161 dep_index,
162 Edge::Dev(
163 group,
164 Some(RequestedExtras::Dependency(&dep.extra)),
165 dep.complexified_marker,
166 ),
167 );
168
169 if seen.insert((&dep.package_id, None)) {
171 queue.push_back((&dep.package_id, None));
172 }
173 for extra in &dep.extra {
174 if seen.insert((&dep.package_id, Some(extra))) {
175 queue.push_back((&dep.package_id, Some(extra)));
176 }
177 }
178 }
179 }
180
181 {
190 let by_name: FxHashMap<_, Vec<_>> = {
192 lock.packages().iter().fold(
193 FxHashMap::with_capacity_and_hasher(lock.len(), FxBuildHasher),
194 |mut map, package| {
195 map.entry(&package.id.name).or_default().push(package);
196 map
197 },
198 )
199 };
200
201 for requirement in lock.requirements() {
203 for package in by_name.get(&requirement.name).into_iter().flatten() {
204 let marker = if package.fork_markers.is_empty() {
207 requirement.marker
208 } else {
209 let mut combined = MarkerTree::FALSE;
210 for fork_marker in &package.fork_markers {
211 combined.or(fork_marker.pep508());
212 }
213 combined.and(requirement.marker);
214 combined
215 };
216 if marker.is_false() {
217 continue;
218 }
219 if markers.is_some_and(|markers| !marker.evaluate(markers, &[])) {
220 continue;
221 }
222 let index = inverse
224 .entry(&package.id)
225 .or_insert_with(|| graph.add_node(Node::Package(&package.id)));
226
227 graph.add_edge(
229 root,
230 *index,
231 Edge::Prod(
232 Some(RequestedExtras::Requirement(requirement.extras.as_ref())),
233 UniversalMarker::from_combined(marker),
234 ),
235 );
236
237 if seen.insert((&package.id, None)) {
239 queue.push_back((&package.id, None));
240 }
241 for extra in &*requirement.extras {
242 if seen.insert((&package.id, Some(extra))) {
243 queue.push_back((&package.id, Some(extra)));
244 }
245 }
246 }
247 }
248
249 for (group, requirements) in lock.dependency_groups() {
251 if !groups.contains(group) {
252 continue;
253 }
254 for requirement in requirements {
255 for package in by_name.get(&requirement.name).into_iter().flatten() {
256 let marker = if package.fork_markers.is_empty() {
259 requirement.marker
260 } else {
261 let mut combined = MarkerTree::FALSE;
262 for fork_marker in &package.fork_markers {
263 combined.or(fork_marker.pep508());
264 }
265 combined.and(requirement.marker);
266 combined
267 };
268 if marker.is_false() {
269 continue;
270 }
271 if markers.is_some_and(|markers| !marker.evaluate(markers, &[])) {
272 continue;
273 }
274 let index = inverse
276 .entry(&package.id)
277 .or_insert_with(|| graph.add_node(Node::Package(&package.id)));
278
279 graph.add_edge(
281 root,
282 *index,
283 Edge::Dev(
284 group,
285 Some(RequestedExtras::Requirement(requirement.extras.as_ref())),
286 UniversalMarker::from_combined(marker),
287 ),
288 );
289
290 if seen.insert((&package.id, None)) {
292 queue.push_back((&package.id, None));
293 }
294 for extra in &*requirement.extras {
295 if seen.insert((&package.id, Some(extra))) {
296 queue.push_back((&package.id, Some(extra)));
297 }
298 }
299 }
300 }
301 }
302 }
303
304 while let Some((id, extra)) = queue.pop_front() {
306 let index = inverse[&id];
307 let package = lock.find_by_id(id);
308
309 let deps = if let Some(extra) = extra {
310 Either::Left(
311 package
312 .optional_dependencies
313 .get(extra)
314 .into_iter()
315 .flatten(),
316 )
317 } else {
318 Either::Right(package.dependencies.iter())
319 };
320
321 for dep in deps {
322 if prune.contains(&dep.package_id.name) {
323 continue;
324 }
325
326 if markers
327 .is_some_and(|markers| !dep.complexified_marker.evaluate_no_extras(markers))
328 {
329 continue;
330 }
331
332 let dep_index = *inverse
334 .entry(&dep.package_id)
335 .or_insert_with(|| graph.add_node(Node::Package(&dep.package_id)));
336
337 graph.add_edge(
339 index,
340 dep_index,
341 if let Some(extra) = extra {
342 Edge::Optional(
343 extra,
344 Some(RequestedExtras::Dependency(&dep.extra)),
345 dep.complexified_marker,
346 )
347 } else {
348 Edge::Prod(
349 Some(RequestedExtras::Dependency(&dep.extra)),
350 dep.complexified_marker,
351 )
352 },
353 );
354
355 if seen.insert((&dep.package_id, None)) {
357 queue.push_back((&dep.package_id, None));
358 }
359 for extra in &dep.extra {
360 if seen.insert((&dep.package_id, Some(extra))) {
361 queue.push_back((&dep.package_id, Some(extra)));
362 }
363 }
364 }
365 }
366
367 {
369 let mut reachable = graph
370 .node_indices()
371 .filter(|index| match graph[*index] {
372 Node::Package(package_id) => members.contains(package_id),
373 Node::Root => true,
374 })
375 .collect::<FxHashSet<_>>();
376 let mut stack = reachable.iter().copied().collect::<VecDeque<_>>();
377 while let Some(node) = stack.pop_front() {
378 for edge in graph.edges_directed(node, Direction::Outgoing) {
379 if reachable.insert(edge.target()) {
380 stack.push_back(edge.target());
381 }
382 }
383 }
384
385 graph.retain_nodes(|_, index| reachable.contains(&index));
387 }
388
389 if invert {
391 graph.reverse();
392 }
393
394 if !packages.is_empty() {
396 let mut reachable = graph
397 .node_indices()
398 .filter(|index| {
399 let Node::Package(package_id) = graph[*index] else {
400 return false;
401 };
402 packages.contains(&package_id.name)
403 })
404 .collect::<FxHashSet<_>>();
405 let mut stack = reachable.iter().copied().collect::<VecDeque<_>>();
406 while let Some(node) = stack.pop_front() {
407 for edge in graph.edges_directed(node, Direction::Outgoing) {
408 if reachable.insert(edge.target()) {
409 stack.push_back(edge.target());
410 }
411 }
412 }
413
414 graph.retain_nodes(|_, index| reachable.contains(&index));
416 }
417
418 let roots = {
420 if !packages.is_empty() {
422 let mut roots = graph
423 .node_indices()
424 .filter(|index| {
425 let Node::Package(package_id) = graph[*index] else {
426 return false;
427 };
428 packages.contains(&package_id.name)
429 })
430 .collect::<Vec<_>>();
431
432 roots.sort_by_key(|index| &graph[*index]);
434
435 roots
436 } else {
437 let mut roots = if invert {
438 graph
441 .node_indices()
442 .filter(|index| {
443 graph
444 .edges_directed(*index, Direction::Incoming)
445 .next()
446 .is_none()
447 })
448 .collect::<Vec<_>>()
449 } else {
450 graph
452 .node_indices()
453 .filter(|index| matches!(graph[*index], Node::Root))
454 .collect::<Vec<_>>()
455 };
456
457 roots.sort_by_key(|index| &graph[*index]);
458 roots
459 }
460 };
461
462 Self {
463 graph,
464 roots,
465 latest,
466 depth,
467 no_dedupe,
468 invert,
469 lock,
470 show_sizes,
471 conflict_marker,
472 }
473 }
474
475 fn visit(
477 &'env self,
478 cursor: Cursor,
479 visited: &mut FxHashMap<VisitedNode<'env>, Vec<&'env PackageId>>,
480 path: &mut Vec<VisitedNode<'env>>,
481 ) -> Vec<String> {
482 if path.len() > self.depth {
484 return Vec::new();
485 }
486
487 let Node::Package(package_id) = self.graph[cursor.node()] else {
488 return Vec::new();
489 };
490 let edge = cursor.edge().map(|edge_id| &self.graph[edge_id]);
491 let package = self.lock.find_by_id(package_id);
492
493 let expanded_extras = self.expanded_extras(package, edge);
494 let visited_node = VisitedNode {
495 package_id,
496 expanded_extras: expanded_extras.clone(),
497 marker: self.invert.then_some(cursor.marker()),
498 };
499
500 let line = {
501 let mut line = format!("{}", package_id.name);
502
503 if let Some(extras) = edge.and_then(Edge::extras) {
504 if !extras.is_empty() {
505 line.push('[');
506 line.push_str(extras.iter().join(", ").as_str());
507 line.push(']');
508 }
509 }
510
511 if let Some(version) = package_id.version.as_ref() {
512 line.push(' ');
513 line.push('v');
514 let _ = write!(line, "{version}");
515 }
516
517 if let Some(edge) = edge {
518 match edge {
519 Edge::Prod(..) => {}
520 Edge::Optional(extra, ..) => {
521 let _ = write!(line, " (extra: {extra})");
522 }
523 Edge::Dev(group, ..) => {
524 let _ = write!(line, " (group: {group})");
525 }
526 }
527 }
528
529 if self.show_sizes {
532 if let Some(size_bytes) = package.wheels.iter().find_map(|wheel| wheel.size) {
533 let (bytes, unit) = human_readable_bytes(size_bytes);
534 line.push(' ');
535 line.push_str(format!("{}", format!("({bytes:.1}{unit})").dimmed()).as_str());
536 }
537 }
538
539 line
540 };
541
542 if path.contains(&visited_node) {
546 return vec![format!("{line} (*)")];
547 }
548 if !self.no_dedupe
549 && let Some(requirements) = visited.get(&visited_node)
550 {
551 return if requirements.is_empty() {
552 vec![line]
553 } else {
554 vec![format!("{line} (*)")]
555 };
556 }
557
558 let line = if let Some(version) = self.latest.get(package_id) {
560 format!("{line} {}", format!("(latest: v{version})").bold().cyan())
561 } else {
562 line
563 };
564
565 let mut dependencies = if self.invert && edge.is_some_and(Edge::is_dev) {
566 Vec::new()
569 } else {
570 self.graph
571 .edges_directed(cursor.node(), Direction::Outgoing)
572 .filter_map(|edge| match self.graph[edge.target()] {
573 Node::Root => None,
574 Node::Package(_) => {
575 let edge_kind = &self.graph[edge.id()];
576
577 if self.invert {
578 if !expanded_extras.is_empty()
581 && edge_kind.extras().is_none_or(|extras| {
582 !expanded_extras.iter().all(|extra| extras.contains(extra))
583 })
584 {
585 return None;
586 }
587
588 let mut marker = cursor.marker();
591 marker.and(edge_kind.marker());
592 if marker.is_false() {
593 return None;
594 }
595 Some(Cursor::new(edge.target(), edge.id(), marker))
596 } else {
597 if let Edge::Optional(required_extra, ..) = edge_kind
600 && !expanded_extras.contains(required_extra)
601 {
602 return None;
603 }
604 Some(Cursor::new(edge.target(), edge.id(), UniversalMarker::TRUE))
605 }
606 }
607 })
608 .collect::<Vec<_>>()
609 };
610 dependencies.sort_by_key(|cursor| {
611 let node = &self.graph[cursor.node()];
612 let edge = cursor
613 .edge()
614 .map(|edge_id| &self.graph[edge_id])
615 .map(Edge::kind);
616 (edge, node)
617 });
618
619 let mut lines = vec![line];
620
621 if path.len() < self.depth {
624 visited.insert(
625 visited_node.clone(),
626 dependencies
627 .iter()
628 .filter_map(|node| match self.graph[node.node()] {
629 Node::Package(package_id) => Some(package_id),
630 Node::Root => None,
631 })
632 .collect(),
633 );
634 }
635 path.push(visited_node);
636
637 for (index, dep) in dependencies.iter().enumerate() {
638 let (prefix_top, prefix_rest) = if dependencies.len() - 1 == index {
658 ("└── ", " ")
659 } else {
660 ("├── ", "│ ")
661 };
662 for (visited_index, visited_line) in self.visit(*dep, visited, path).iter().enumerate()
663 {
664 let prefix = if visited_index == 0 {
665 prefix_top
666 } else {
667 prefix_rest
668 };
669 lines.push(format!("{prefix}{visited_line}"));
670 }
671 }
672
673 path.pop();
674
675 lines
676 }
677
678 fn render(&self) -> Vec<String> {
680 let mut path = Vec::new();
681 let mut lines = Vec::with_capacity(self.graph.node_count());
682 let mut visited =
683 FxHashMap::with_capacity_and_hasher(self.graph.node_count(), FxBuildHasher);
684
685 for node in &self.roots {
686 match self.graph[*node] {
687 Node::Root => {
688 for edge in self.graph.edges_directed(*node, Direction::Outgoing) {
689 let node = edge.target();
690 path.clear();
691 lines.extend(self.visit(
692 Cursor::new(node, edge.id(), self.conflict_marker),
693 &mut visited,
694 &mut path,
695 ));
696 }
697 }
698 Node::Package(_) => {
699 path.clear();
700 lines.extend(self.visit(
701 Cursor::root(*node, self.conflict_marker),
702 &mut visited,
703 &mut path,
704 ));
705 }
706 }
707 }
708
709 lines
710 }
711
712 fn expanded_extras(
714 &self,
715 package: &'env Package,
716 edge: Option<&Edge<'env>>,
717 ) -> BTreeSet<&'env ExtraName> {
718 if self.invert {
719 return edge.and_then(Edge::required_extra).into_iter().collect();
722 }
723
724 let Some(requested_extras) = edge.and_then(Edge::extras) else {
725 return package.optional_dependencies.keys().collect();
727 };
728
729 requested_extras
730 .iter()
731 .filter(|extra| package.optional_dependencies.contains_key(*extra))
732 .collect()
733 }
734}
735
736#[derive(Debug, Clone, PartialEq, Eq, Hash)]
737struct VisitedNode<'env> {
738 package_id: &'env PackageId,
739 expanded_extras: BTreeSet<&'env ExtraName>,
740 marker: Option<UniversalMarker>,
741}
742
743#[derive(Debug, Clone, PartialEq, Eq, Ord, PartialOrd)]
744enum Node<'env> {
745 Root,
747 Package(&'env PackageId),
749}
750
751#[derive(Debug, Clone, PartialEq, Eq, Ord, PartialOrd)]
752enum Edge<'env> {
753 Prod(Option<RequestedExtras<'env>>, UniversalMarker),
754 Optional(
755 &'env ExtraName,
756 Option<RequestedExtras<'env>>,
757 UniversalMarker,
758 ),
759 Dev(
760 &'env GroupName,
761 Option<RequestedExtras<'env>>,
762 UniversalMarker,
763 ),
764}
765
766impl<'env> Edge<'env> {
767 fn extras(&self) -> Option<RequestedExtras<'env>> {
768 match self {
769 Self::Prod(extras, _) => *extras,
770 Self::Optional(_, extras, _) => *extras,
771 Self::Dev(_, extras, _) => *extras,
772 }
773 }
774
775 fn required_extra(&self) -> Option<&'env ExtraName> {
776 match self {
777 Self::Optional(extra, ..) => Some(extra),
778 Self::Prod(..) | Self::Dev(..) => None,
779 }
780 }
781
782 fn marker(&self) -> UniversalMarker {
783 match self {
784 Self::Prod(_, marker) | Self::Optional(_, _, marker) | Self::Dev(_, _, marker) => {
785 *marker
786 }
787 }
788 }
789
790 fn is_dev(&self) -> bool {
791 matches!(self, Self::Dev(..))
792 }
793
794 fn kind(&self) -> EdgeKind<'env> {
795 match self {
796 Self::Prod(..) => EdgeKind::Prod,
797 Self::Optional(extra, ..) => EdgeKind::Optional(extra),
798 Self::Dev(group, ..) => EdgeKind::Dev(group),
799 }
800 }
801}
802
803#[derive(Debug, Copy, Clone)]
804enum RequestedExtras<'env> {
805 Dependency(&'env BTreeSet<ExtraName>),
806 Requirement(&'env [ExtraName]),
807}
808
809impl PartialEq for RequestedExtras<'_> {
810 fn eq(&self, other: &Self) -> bool {
811 self.iter().eq(other.iter())
812 }
813}
814
815impl Eq for RequestedExtras<'_> {}
816
817impl PartialOrd for RequestedExtras<'_> {
818 fn partial_cmp(&self, other: &Self) -> Option<Ordering> {
819 Some(self.cmp(other))
820 }
821}
822
823impl Ord for RequestedExtras<'_> {
824 fn cmp(&self, other: &Self) -> Ordering {
825 self.iter().cmp(other.iter())
826 }
827}
828
829impl<'env> RequestedExtras<'env> {
830 fn contains(self, extra: &ExtraName) -> bool {
831 match self {
832 Self::Dependency(extras) => extras.contains(extra),
833 Self::Requirement(extras) => extras.contains(extra),
834 }
835 }
836
837 fn is_empty(self) -> bool {
838 match self {
839 Self::Dependency(extras) => extras.is_empty(),
840 Self::Requirement(extras) => extras.is_empty(),
841 }
842 }
843
844 fn iter(self) -> impl Iterator<Item = &'env ExtraName> {
845 match self {
846 Self::Dependency(extras) => Either::Left(extras.iter()),
847 Self::Requirement(extras) => Either::Right(extras.iter()),
848 }
849 }
850}
851
852#[derive(Debug, Clone, PartialEq, Eq, Ord, PartialOrd)]
853enum EdgeKind<'env> {
854 Prod,
855 Optional(&'env ExtraName),
856 Dev(&'env GroupName),
857}
858
859#[derive(Debug, Copy, Clone, PartialEq, Eq, Ord, PartialOrd)]
861struct Cursor(NodeIndex, Option<EdgeIndex>, UniversalMarker);
862
863impl Cursor {
864 fn new(node: NodeIndex, edge: EdgeIndex, marker: UniversalMarker) -> Self {
866 Self(node, Some(edge), marker)
867 }
868
869 fn root(node: NodeIndex, marker: UniversalMarker) -> Self {
871 Self(node, None, marker)
872 }
873
874 fn node(&self) -> NodeIndex {
876 self.0
877 }
878
879 fn edge(&self) -> Option<EdgeIndex> {
881 self.1
882 }
883
884 fn marker(&self) -> UniversalMarker {
886 self.2
887 }
888}
889
890impl std::fmt::Display for TreeDisplay<'_> {
891 fn fmt(&self, f: &mut std::fmt::Formatter) -> std::fmt::Result {
892 use owo_colors::OwoColorize;
893
894 let mut deduped = false;
895 for line in self.render() {
896 deduped |= line.contains('*');
897 writeln!(f, "{line}")?;
898 }
899
900 if deduped {
901 let message = if self.no_dedupe {
902 "(*) Package tree is a cycle and cannot be shown".italic()
903 } else {
904 "(*) Package tree already displayed".italic()
905 };
906 writeln!(f, "{message}")?;
907 }
908
909 Ok(())
910 }
911}