1use std::cmp::Ordering;
2use std::collections::{BTreeMap, BTreeSet, VecDeque};
3use std::fmt::Write;
4use std::path::Path;
5
6use either::Either;
7use itertools::Itertools;
8use owo_colors::OwoColorize;
9use petgraph::graph::{EdgeIndex, NodeIndex};
10use petgraph::prelude::EdgeRef;
11use petgraph::{Direction, Graph};
12use rustc_hash::{FxBuildHasher, FxHashMap, FxHashSet};
13use serde::Serialize;
14
15use uv_configuration::DependencyGroupsWithDefaults;
16use uv_console::human_readable_bytes;
17use uv_fs::PortablePathBuf;
18use uv_normalize::{ExtraName, GroupName, PackageName};
19use uv_pep440::Version;
20use uv_pep508::MarkerTree;
21use uv_pypi_types::ResolverMarkerEnvironment;
22
23use crate::lock::export::{
24 MetadataNode, MetadataNodeId, MetadataNodeKind, MetadataScript, MetadataWorkspace,
25 MetadataWorkspaceMember,
26};
27use crate::lock::{Package, PackageId};
28use crate::{ConflictMarker, Lock, PackageMap, UniversalMarker};
29
30#[derive(Debug, Clone, Copy)]
31pub enum TreeJsonTarget<'a> {
32 Workspace(&'a Path),
33 Script(&'a Path),
34}
35
36impl<'a> TreeJsonTarget<'a> {
37 fn root(self) -> &'a Path {
38 match self {
39 Self::Workspace(root) => root,
40 Self::Script(script) => script.parent().unwrap_or_else(|| Path::new("")),
41 }
42 }
43}
44
45#[derive(Debug)]
46pub struct TreeDisplay<'env> {
47 graph: petgraph::graph::Graph<Node<'env>, Edge<'env>, petgraph::Directed>,
49 roots: Vec<NodeIndex>,
51 latest: &'env PackageMap<Version>,
53 depth: usize,
55 no_dedupe: bool,
57 invert: bool,
59 prod: bool,
61 groups: DependencyGroupsWithDefaults,
63 lock: &'env Lock,
65 show_sizes: bool,
67 conflict_marker: UniversalMarker,
69}
70
71impl<'env> TreeDisplay<'env> {
72 pub fn new(
74 lock: &'env Lock,
75 markers: Option<&'env ResolverMarkerEnvironment>,
76 latest: &'env PackageMap<Version>,
77 depth: usize,
78 prune: &[PackageName],
79 packages: &[PackageName],
80 groups: &DependencyGroupsWithDefaults,
81 no_dedupe: bool,
82 invert: bool,
83 show_sizes: bool,
84 ) -> Self {
85 let members: BTreeSet<&PackageId> = if lock.members().is_empty() {
93 lock.root().into_iter().map(|package| &package.id).collect()
94 } else {
95 lock.packages
96 .iter()
97 .filter_map(|package| {
98 if lock.members().contains(&package.id.name) {
99 Some(&package.id)
100 } else {
101 None
102 }
103 })
104 .collect()
105 };
106
107 let conflict_marker = UniversalMarker::new(
110 MarkerTree::TRUE,
111 ConflictMarker::from_conflicts(lock.conflicts()),
112 );
113
114 let size_guess = lock.packages.len();
116 let mut graph =
117 Graph::<Node, Edge, petgraph::Directed>::with_capacity(size_guess, size_guess);
118 let mut inverse = FxHashMap::with_capacity_and_hasher(size_guess, FxBuildHasher);
119 let mut queue: VecDeque<(&PackageId, Option<&ExtraName>)> = VecDeque::new();
120 let mut seen = FxHashSet::default();
121
122 let root = graph.add_node(Node::Root);
123
124 for id in members.iter().copied() {
126 if prune.contains(&id.name) {
127 continue;
128 }
129
130 let dist = lock.find_by_id(id);
131
132 let index = *inverse
136 .entry(id)
137 .or_insert_with(|| graph.add_node(Node::Package(id)));
138
139 graph.add_edge(root, index, Edge::Prod(None, UniversalMarker::TRUE));
141
142 if groups.prod() {
143 if seen.insert((id, None)) {
145 queue.push_back((id, None));
146 }
147
148 for extra in dist.optional_dependencies.keys() {
150 if seen.insert((id, Some(extra))) {
151 queue.push_back((id, Some(extra)));
152 }
153 }
154 }
155
156 for (group, dep) in dist
158 .dependency_groups
159 .iter()
160 .filter_map(|(group, deps)| {
161 if groups.contains(group) {
162 Some(deps.iter().map(move |dep| (group, dep)))
163 } else {
164 None
165 }
166 })
167 .flatten()
168 {
169 if prune.contains(&dep.package_id.name) {
170 continue;
171 }
172
173 if markers
174 .is_some_and(|markers| !dep.complexified_marker.evaluate_no_extras(markers))
175 {
176 continue;
177 }
178
179 let dep_index = *inverse
181 .entry(&dep.package_id)
182 .or_insert_with(|| graph.add_node(Node::Package(&dep.package_id)));
183
184 graph.add_edge(
186 index,
187 dep_index,
188 Edge::Dev(
189 group,
190 Some(RequestedExtras::Dependency(&dep.extra)),
191 dep.complexified_marker,
192 ),
193 );
194
195 if seen.insert((&dep.package_id, None)) {
197 queue.push_back((&dep.package_id, None));
198 }
199 for extra in &dep.extra {
200 if seen.insert((&dep.package_id, Some(extra))) {
201 queue.push_back((&dep.package_id, Some(extra)));
202 }
203 }
204 }
205 }
206
207 {
216 let by_name: FxHashMap<_, Vec<_>> = {
218 lock.packages().iter().fold(
219 FxHashMap::with_capacity_and_hasher(lock.len(), FxBuildHasher),
220 |mut map, package| {
221 map.entry(&package.id.name).or_default().push(package);
222 map
223 },
224 )
225 };
226
227 for requirement in lock.requirements() {
229 for package in by_name.get(&requirement.name).into_iter().flatten() {
230 let marker = if package.fork_markers.is_empty() {
233 requirement.marker
234 } else {
235 let mut combined = MarkerTree::FALSE;
236 for fork_marker in &package.fork_markers {
237 combined.or(fork_marker.pep508());
238 }
239 combined.and(requirement.marker);
240 combined
241 };
242 if marker.is_false() {
243 continue;
244 }
245 if markers.is_some_and(|markers| !marker.evaluate(markers, &[])) {
246 continue;
247 }
248 let index = inverse
250 .entry(&package.id)
251 .or_insert_with(|| graph.add_node(Node::Package(&package.id)));
252
253 graph.add_edge(
255 root,
256 *index,
257 Edge::Prod(
258 Some(RequestedExtras::Requirement(requirement.extras.as_ref())),
259 UniversalMarker::from_combined(marker),
260 ),
261 );
262
263 if seen.insert((&package.id, None)) {
265 queue.push_back((&package.id, None));
266 }
267 for extra in &*requirement.extras {
268 if seen.insert((&package.id, Some(extra))) {
269 queue.push_back((&package.id, Some(extra)));
270 }
271 }
272 }
273 }
274
275 for (group, requirements) in lock.dependency_groups() {
277 if !groups.contains(group) {
278 continue;
279 }
280 for requirement in requirements {
281 for package in by_name.get(&requirement.name).into_iter().flatten() {
282 let marker = if package.fork_markers.is_empty() {
285 requirement.marker
286 } else {
287 let mut combined = MarkerTree::FALSE;
288 for fork_marker in &package.fork_markers {
289 combined.or(fork_marker.pep508());
290 }
291 combined.and(requirement.marker);
292 combined
293 };
294 if marker.is_false() {
295 continue;
296 }
297 if markers.is_some_and(|markers| !marker.evaluate(markers, &[])) {
298 continue;
299 }
300 let index = inverse
302 .entry(&package.id)
303 .or_insert_with(|| graph.add_node(Node::Package(&package.id)));
304
305 graph.add_edge(
307 root,
308 *index,
309 Edge::Dev(
310 group,
311 Some(RequestedExtras::Requirement(requirement.extras.as_ref())),
312 UniversalMarker::from_combined(marker),
313 ),
314 );
315
316 if seen.insert((&package.id, None)) {
318 queue.push_back((&package.id, None));
319 }
320 for extra in &*requirement.extras {
321 if seen.insert((&package.id, Some(extra))) {
322 queue.push_back((&package.id, Some(extra)));
323 }
324 }
325 }
326 }
327 }
328 }
329
330 while let Some((id, extra)) = queue.pop_front() {
332 let index = inverse[&id];
333 let package = lock.find_by_id(id);
334
335 let deps = if let Some(extra) = extra {
336 Either::Left(
337 package
338 .optional_dependencies
339 .get(extra)
340 .into_iter()
341 .flatten(),
342 )
343 } else {
344 Either::Right(package.dependencies.iter())
345 };
346
347 for dep in deps {
348 if prune.contains(&dep.package_id.name) {
349 continue;
350 }
351
352 if markers
353 .is_some_and(|markers| !dep.complexified_marker.evaluate_no_extras(markers))
354 {
355 continue;
356 }
357
358 let dep_index = *inverse
360 .entry(&dep.package_id)
361 .or_insert_with(|| graph.add_node(Node::Package(&dep.package_id)));
362
363 graph.add_edge(
365 index,
366 dep_index,
367 if let Some(extra) = extra {
368 Edge::Optional(
369 extra,
370 Some(RequestedExtras::Dependency(&dep.extra)),
371 dep.complexified_marker,
372 )
373 } else {
374 Edge::Prod(
375 Some(RequestedExtras::Dependency(&dep.extra)),
376 dep.complexified_marker,
377 )
378 },
379 );
380
381 if seen.insert((&dep.package_id, None)) {
383 queue.push_back((&dep.package_id, None));
384 }
385 for extra in &dep.extra {
386 if seen.insert((&dep.package_id, Some(extra))) {
387 queue.push_back((&dep.package_id, Some(extra)));
388 }
389 }
390 }
391 }
392
393 {
395 let mut reachable = graph
396 .node_indices()
397 .filter(|index| match graph[*index] {
398 Node::Package(package_id) => members.contains(package_id),
399 Node::Root => true,
400 })
401 .collect::<FxHashSet<_>>();
402 let mut stack = reachable.iter().copied().collect::<VecDeque<_>>();
403 while let Some(node) = stack.pop_front() {
404 for edge in graph.edges_directed(node, Direction::Outgoing) {
405 if reachable.insert(edge.target()) {
406 stack.push_back(edge.target());
407 }
408 }
409 }
410
411 graph.retain_nodes(|_, index| reachable.contains(&index));
413 }
414
415 if invert {
417 graph.reverse();
418 }
419
420 if !packages.is_empty() {
422 let mut reachable = 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::<FxHashSet<_>>();
431 let mut stack = reachable.iter().copied().collect::<VecDeque<_>>();
432 while let Some(node) = stack.pop_front() {
433 for edge in graph.edges_directed(node, Direction::Outgoing) {
434 if reachable.insert(edge.target()) {
435 stack.push_back(edge.target());
436 }
437 }
438 }
439
440 graph.retain_nodes(|_, index| reachable.contains(&index));
442 }
443
444 let roots = {
446 if !packages.is_empty() {
448 let mut roots = graph
449 .node_indices()
450 .filter(|index| {
451 let Node::Package(package_id) = graph[*index] else {
452 return false;
453 };
454 packages.contains(&package_id.name)
455 })
456 .collect::<Vec<_>>();
457
458 roots.sort_by_key(|index| &graph[*index]);
460
461 roots
462 } else {
463 let mut roots = if invert {
464 graph
467 .node_indices()
468 .filter(|index| {
469 graph
470 .edges_directed(*index, Direction::Incoming)
471 .next()
472 .is_none()
473 })
474 .collect::<Vec<_>>()
475 } else {
476 graph
478 .node_indices()
479 .filter(|index| matches!(graph[*index], Node::Root))
480 .collect::<Vec<_>>()
481 };
482
483 roots.sort_by_key(|index| &graph[*index]);
484 roots
485 }
486 };
487
488 Self {
489 graph,
490 roots,
491 latest,
492 depth,
493 no_dedupe,
494 invert,
495 prod: groups.prod(),
496 groups: groups.clone(),
497 lock,
498 show_sizes,
499 conflict_marker,
500 }
501 }
502
503 fn visit(
505 &'env self,
506 cursor: Cursor,
507 visited: &mut FxHashMap<VisitedNode<'env>, Vec<&'env PackageId>>,
508 path: &mut Vec<VisitedNode<'env>>,
509 ) -> Vec<String> {
510 if path.len() > self.depth {
512 return Vec::new();
513 }
514
515 let Node::Package(package_id) = self.graph[cursor.node()] else {
516 return Vec::new();
517 };
518 let edge = cursor.edge().map(|edge_id| &self.graph[edge_id]);
519 let package = self.lock.find_by_id(package_id);
520
521 let expanded_extras = self.expanded_extras(package, edge);
522 let visited_node = VisitedNode {
523 package_id,
524 expanded_extras: expanded_extras.clone(),
525 marker: self.invert.then_some(cursor.marker()),
526 };
527
528 let line = {
529 let mut line = format!("{}", package_id.name);
530
531 if let Some(extras) = edge.and_then(Edge::extras) {
532 if !extras.is_empty() {
533 line.push('[');
534 line.push_str(extras.iter().join(", ").as_str());
535 line.push(']');
536 }
537 }
538
539 if let Some(version) = package_id.version.as_ref() {
540 line.push(' ');
541 line.push('v');
542 let _ = write!(line, "{version}");
543 }
544
545 if let Some(edge) = edge {
546 match edge {
547 Edge::Prod(..) => {}
548 Edge::Optional(extra, ..) => {
549 let _ = write!(line, " (extra: {extra})");
550 }
551 Edge::Dev(group, ..) => {
552 let _ = write!(line, " (group: {group})");
553 }
554 }
555 }
556
557 if self.show_sizes {
560 if let Some(size_bytes) = package.wheels.iter().find_map(|wheel| wheel.size) {
561 let (bytes, unit) = human_readable_bytes(size_bytes);
562 line.push(' ');
563 line.push_str(format!("{}", format!("({bytes:.1}{unit})").dimmed()).as_str());
564 }
565 }
566
567 line
568 };
569
570 if path.contains(&visited_node) {
574 return vec![format!("{line} (*)")];
575 }
576 if !self.no_dedupe
577 && let Some(requirements) = visited.get(&visited_node)
578 {
579 return if requirements.is_empty() {
580 vec![line]
581 } else {
582 vec![format!("{line} (*)")]
583 };
584 }
585
586 let line = if let Some(version) = self.latest.get(package_id) {
588 format!("{line} {}", format!("(latest: v{version})").bold().cyan())
589 } else {
590 line
591 };
592
593 let mut dependencies = if self.invert && edge.is_some_and(Edge::is_dev) {
594 Vec::new()
597 } else {
598 self.graph
599 .edges_directed(cursor.node(), Direction::Outgoing)
600 .filter_map(|edge| match self.graph[edge.target()] {
601 Node::Root => None,
602 Node::Package(_) => {
603 let edge_kind = &self.graph[edge.id()];
604
605 if self.invert {
606 if !expanded_extras.is_empty()
609 && edge_kind.extras().is_none_or(|extras| {
610 !expanded_extras.iter().all(|extra| extras.contains(extra))
611 })
612 {
613 return None;
614 }
615
616 let mut marker = cursor.marker();
619 marker.and(edge_kind.marker());
620 if marker.is_false() {
621 return None;
622 }
623 Some(Cursor::new(edge.target(), edge.id(), marker))
624 } else {
625 if let Edge::Optional(required_extra, ..) = edge_kind
628 && !expanded_extras.contains(required_extra)
629 {
630 return None;
631 }
632 Some(Cursor::new(edge.target(), edge.id(), UniversalMarker::TRUE))
633 }
634 }
635 })
636 .collect::<Vec<_>>()
637 };
638 dependencies.sort_by_key(|cursor| {
639 let node = &self.graph[cursor.node()];
640 let edge = cursor
641 .edge()
642 .map(|edge_id| &self.graph[edge_id])
643 .map(Edge::kind);
644 (edge, node)
645 });
646
647 let mut lines = vec![line];
648
649 if path.len() < self.depth {
652 visited.insert(
653 visited_node.clone(),
654 dependencies
655 .iter()
656 .filter_map(|node| match self.graph[node.node()] {
657 Node::Package(package_id) => Some(package_id),
658 Node::Root => None,
659 })
660 .collect(),
661 );
662 }
663 path.push(visited_node);
664
665 for (index, dep) in dependencies.iter().enumerate() {
666 let (prefix_top, prefix_rest) = if dependencies.len() - 1 == index {
686 ("└── ", " ")
687 } else {
688 ("├── ", "│ ")
689 };
690 for (visited_index, visited_line) in self.visit(*dep, visited, path).iter().enumerate()
691 {
692 let prefix = if visited_index == 0 {
693 prefix_top
694 } else {
695 prefix_rest
696 };
697 lines.push(format!("{prefix}{visited_line}"));
698 }
699 }
700
701 path.pop();
702
703 lines
704 }
705
706 fn render(&self) -> Vec<String> {
708 let mut path = Vec::new();
709 let mut lines = Vec::with_capacity(self.graph.node_count());
710 let mut visited =
711 FxHashMap::with_capacity_and_hasher(self.graph.node_count(), FxBuildHasher);
712
713 for node in &self.roots {
714 match self.graph[*node] {
715 Node::Root => {
716 for edge in self.graph.edges_directed(*node, Direction::Outgoing) {
717 let node = edge.target();
718 path.clear();
719 lines.extend(self.visit(
720 Cursor::new(node, edge.id(), self.conflict_marker),
721 &mut visited,
722 &mut path,
723 ));
724 }
725 }
726 Node::Package(_) => {
727 path.clear();
728 lines.extend(self.visit(
729 Cursor::root(*node, self.conflict_marker),
730 &mut visited,
731 &mut path,
732 ));
733 }
734 }
735 }
736
737 lines
738 }
739
740 fn expanded_extras(
742 &self,
743 package: &'env Package,
744 edge: Option<&Edge<'env>>,
745 ) -> BTreeSet<&'env ExtraName> {
746 if self.invert {
747 return edge.and_then(Edge::required_extra).into_iter().collect();
750 }
751
752 let Some(requested_extras) = edge.and_then(Edge::extras) else {
753 return package.optional_dependencies.keys().collect();
755 };
756
757 requested_extras
758 .iter()
759 .filter(|extra| package.optional_dependencies.contains_key(*extra))
760 .collect()
761 }
762
763 pub fn to_json(&self, target: TreeJsonTarget<'_>) -> Result<String, serde_json::Error> {
765 serde_json::to_string_pretty(&JsonGraph::new(self, target))
766 }
767
768 fn json_traversal(&self) -> JsonTraversal {
778 let mut distances = FxHashMap::default();
779 let mut queue = VecDeque::new();
780 let mut nodes = FxHashSet::default();
781 let mut edges = FxHashSet::default();
782
783 for root in &self.roots {
784 match self.graph[*root] {
785 Node::Root => {
786 for edge in self.graph.edges_directed(*root, Direction::Outgoing) {
787 let Node::Package(package_id) = self.graph[edge.target()] else {
788 continue;
789 };
790 let state = JsonTraversalNode {
791 index: edge.target(),
792 expanded_extras: self.expanded_extras(
793 self.lock.find_by_id(package_id),
794 Some(edge.weight()),
795 ),
796 marker: UniversalMarker::TRUE,
797 reached_via_dependency_group: false,
798 };
799 nodes.insert(state.index);
800 if distances.insert(state.clone(), 0).is_none() {
801 queue.push_back(state);
802 }
803 }
804 }
805 Node::Package(package_id) => {
806 let state = JsonTraversalNode {
807 index: *root,
808 expanded_extras: self
809 .expanded_extras(self.lock.find_by_id(package_id), None),
810 marker: if self.invert {
811 self.conflict_marker
812 } else {
813 UniversalMarker::TRUE
814 },
815 reached_via_dependency_group: false,
816 };
817 nodes.insert(state.index);
818 if distances.insert(state.clone(), 0).is_none() {
819 queue.push_back(state);
820 }
821 }
822 }
823 }
824
825 while let Some(source) = queue.pop_front() {
826 let distance = distances[&source];
827 if distance >= self.depth || self.invert && source.reached_via_dependency_group {
828 continue;
829 }
830
831 for edge in self.graph.edges_directed(source.index, Direction::Outgoing) {
832 let edge_kind = edge.weight();
833 let marker = if self.invert {
834 if !source.expanded_extras.is_empty()
837 && edge_kind.extras().is_none_or(|extras| {
838 !source
839 .expanded_extras
840 .iter()
841 .all(|extra| extras.contains(extra))
842 })
843 {
844 continue;
845 }
846
847 let mut marker = source.marker;
850 marker.and(edge_kind.marker());
851 if marker.is_false() {
852 continue;
853 }
854 marker
855 } else {
856 if let Edge::Optional(required_extra, ..) = edge_kind
859 && !source.expanded_extras.contains(required_extra)
860 {
861 continue;
862 }
863 UniversalMarker::TRUE
864 };
865
866 let target = edge.target();
867 if matches!(self.graph[target], Node::Root) {
868 edges.insert(edge.id());
869 continue;
870 }
871 let Node::Package(package_id) = self.graph[target] else {
872 continue;
873 };
874 let state = JsonTraversalNode {
875 index: target,
876 expanded_extras: self
877 .expanded_extras(self.lock.find_by_id(package_id), Some(edge.weight())),
878 marker,
879 reached_via_dependency_group: self.invert && edge_kind.is_dev(),
880 };
881 nodes.insert(state.index);
882 edges.insert(edge.id());
883 if !distances.contains_key(&state) {
884 distances.insert(state.clone(), distance + 1);
885 queue.push_back(state);
886 }
887 }
888 }
889
890 JsonTraversal { nodes, edges }
891 }
892}
893
894#[derive(Debug)]
895struct JsonTraversal {
896 nodes: FxHashSet<NodeIndex>,
897 edges: FxHashSet<EdgeIndex>,
898}
899
900#[derive(Debug, Clone, PartialEq, Eq, Hash)]
901struct JsonTraversalNode<'env> {
902 index: NodeIndex,
903 expanded_extras: BTreeSet<&'env ExtraName>,
904 marker: UniversalMarker,
905 reached_via_dependency_group: bool,
906}
907
908#[derive(Debug, Serialize)]
1008struct JsonGraph {
1009 schema: JsonSchema,
1010 workspace_root: PortablePathBuf,
1011 #[serde(skip_serializing_if = "Option::is_none")]
1012 script: Option<MetadataScript>,
1013 #[serde(skip_serializing_if = "Option::is_none")]
1014 workspace: Option<MetadataWorkspace>,
1015 roots: Vec<JsonRoot>,
1016 inverted: bool,
1017 #[serde(skip_serializing_if = "Vec::is_empty")]
1019 members: Vec<MetadataWorkspaceMember>,
1020 resolution: BTreeMap<String, MetadataNode>,
1021}
1022
1023impl JsonGraph {
1024 fn new(tree: &TreeDisplay<'_>, target: TreeJsonTarget<'_>) -> Self {
1025 let traversal = tree.json_traversal();
1026 let workspace_root = PortablePathBuf::from(target.root());
1027 let mut builder = JsonGraphBuilder::new(tree, workspace_root.clone());
1028
1029 for index in traversal.nodes.iter().copied() {
1030 let Node::Package(package_id) = tree.graph[index] else {
1031 continue;
1032 };
1033 builder.ensure_package(package_id, MetadataNodeKind::Package);
1034 }
1035
1036 for edge in tree
1037 .graph
1038 .edge_references()
1039 .filter(|edge| traversal.edges.contains(&edge.id()))
1040 {
1041 builder.add_package_edge(edge.source(), edge.target(), edge.weight());
1042 }
1043
1044 builder.add_target_edges(target, &traversal);
1045 let (script, workspace) = match target {
1046 TreeJsonTarget::Script(path) => {
1047 let path = PortablePathBuf::from(path);
1048 let id = builder.ensure_script(path.as_ref());
1049 (Some(MetadataScript::new(path, id)), None)
1050 }
1051 TreeJsonTarget::Workspace(path) => {
1052 let path = PortablePathBuf::from(path);
1053 let id = builder.ensure_workspace();
1054 (None, Some(MetadataWorkspace::new(path, id)))
1055 }
1056 };
1057 let roots = builder.roots(target);
1058 let members = builder.members(target);
1059 let resolution = builder.finish();
1060
1061 Self {
1062 schema: JsonSchema {
1063 version: JsonSchemaVersion::Preview,
1064 },
1065 workspace_root,
1066 script,
1067 workspace,
1068 roots,
1069 inverted: tree.invert,
1070 members,
1071 resolution,
1072 }
1073 }
1074}
1075
1076struct JsonGraphBuilder<'tree, 'env> {
1077 tree: &'tree TreeDisplay<'env>,
1078 workspace_root: PortablePathBuf,
1079 resolution: BTreeMap<String, MetadataNode>,
1080}
1081
1082impl<'tree, 'env> JsonGraphBuilder<'tree, 'env> {
1083 fn new(tree: &'tree TreeDisplay<'env>, workspace_root: PortablePathBuf) -> Self {
1084 Self {
1085 tree,
1086 workspace_root,
1087 resolution: BTreeMap::new(),
1088 }
1089 }
1090
1091 fn ensure_node(&mut self, identity: MetadataNodeId) -> String {
1092 let id = identity.to_flat();
1093 self.resolution
1094 .entry(id.clone())
1095 .or_insert_with(|| MetadataNode::new(identity));
1096 id
1097 }
1098
1099 fn ensure_package(&mut self, package_id: &'env PackageId, kind: MetadataNodeKind) -> String {
1100 let is_package = matches!(kind, MetadataNodeKind::Package);
1101 let id = MetadataNodeId::from_package_id(&self.workspace_root, package_id, kind.clone())
1102 .to_flat();
1103 self.resolution.entry(id.clone()).or_insert_with(|| {
1104 let package = self.tree.lock.find_by_id(package_id);
1105 let mut node = MetadataNode::from_package_id(&self.workspace_root, package_id, kind);
1106 if is_package {
1107 node.set_latest_version(self.tree.latest.get(package_id).cloned());
1108 node.set_wheels_from_package(&self.workspace_root, package);
1109 }
1110 node
1111 });
1112 id
1113 }
1114
1115 fn ensure_extra(&mut self, package_id: &'env PackageId, extra: &ExtraName) -> String {
1116 let package = self.ensure_package(package_id, MetadataNodeKind::Package);
1117 let extra_id = self.ensure_package(package_id, MetadataNodeKind::Extra(extra.clone()));
1118 self.add_link(
1119 package.clone(),
1120 extra_id.clone(),
1121 JsonLink::Optional(extra.clone()),
1122 );
1123 self.add_link(extra_id.clone(), package, JsonLink::Dependency(None));
1124 extra_id
1125 }
1126
1127 fn ensure_group(&mut self, package_id: &'env PackageId, group: &GroupName) -> String {
1128 let package = self.ensure_package(package_id, MetadataNodeKind::Package);
1129 let group_id = self.ensure_package(package_id, MetadataNodeKind::Group(group.clone()));
1130 self.add_link(package, group_id.clone(), JsonLink::Group(group.clone()));
1131 group_id
1132 }
1133
1134 fn ensure_workspace(&mut self) -> String {
1135 self.ensure_node(MetadataNodeId::from_workspace(self.workspace_root.clone()))
1136 }
1137
1138 fn ensure_workspace_group(&mut self, group: &GroupName) -> String {
1139 let workspace = self.ensure_workspace();
1140 let group_id = self.ensure_node(MetadataNodeId::from_workspace_group(
1141 self.workspace_root.clone(),
1142 group.clone(),
1143 ));
1144 self.add_link(workspace, group_id.clone(), JsonLink::Group(group.clone()));
1145 group_id
1146 }
1147
1148 fn ensure_script(&mut self, path: &Path) -> String {
1149 self.ensure_node(MetadataNodeId::from_script(PortablePathBuf::from(path)))
1150 }
1151
1152 fn dependency_targets(
1153 &mut self,
1154 package_id: &'env PackageId,
1155 extras: Option<RequestedExtras<'env>>,
1156 ) -> Vec<String> {
1157 let Some(extras) = extras.filter(|extras| !extras.is_empty()) else {
1158 return vec![self.ensure_package(package_id, MetadataNodeKind::Package)];
1159 };
1160 extras
1161 .iter()
1162 .map(|extra| self.ensure_extra(package_id, extra))
1163 .collect()
1164 }
1165
1166 fn add_package_edge(&mut self, source: NodeIndex, target: NodeIndex, edge: &Edge<'env>) {
1167 let (source, target) = if self.tree.invert {
1168 (target, source)
1169 } else {
1170 (source, target)
1171 };
1172 let (Node::Package(source), Node::Package(target)) =
1173 (&self.tree.graph[source], &self.tree.graph[target])
1174 else {
1175 return;
1176 };
1177 let (source, target) = (*source, *target);
1178
1179 let source = match edge {
1180 Edge::Prod(..) => self.ensure_package(source, MetadataNodeKind::Package),
1181 Edge::Optional(extra, ..) => self.ensure_extra(source, extra),
1182 Edge::Dev(group, ..) => self.ensure_group(source, group),
1183 };
1184 let marker = self.marker(edge);
1185 for target in self.dependency_targets(target, edge.extras()) {
1186 self.add_link(source.clone(), target, JsonLink::Dependency(marker.clone()));
1187 }
1188 }
1189
1190 fn add_target_edges(&mut self, target: TreeJsonTarget<'_>, traversal: &JsonTraversal) {
1191 let edges = self
1195 .tree
1196 .graph
1197 .edge_references()
1198 .filter(|edge| !self.tree.invert || traversal.edges.contains(&edge.id()))
1199 .filter_map(|edge| {
1200 let package = match (
1201 &self.tree.graph[edge.source()],
1202 &self.tree.graph[edge.target()],
1203 ) {
1204 (Node::Root, Node::Package(package)) | (Node::Package(package), Node::Root) => {
1205 *package
1206 }
1207 (Node::Root, Node::Root) | (Node::Package(_), Node::Package(_)) => return None,
1208 };
1209 Some((package, edge.weight()))
1210 })
1211 .collect::<Vec<_>>();
1212
1213 match target {
1214 TreeJsonTarget::Script(path) => {
1215 let script = self.ensure_script(path);
1216 for (package, edge) in edges {
1217 let marker = self.marker(edge);
1218 for package in self.dependency_targets(package, edge.extras()) {
1219 self.add_link(
1220 script.clone(),
1221 package,
1222 JsonLink::Dependency(marker.clone()),
1223 );
1224 }
1225 }
1226 }
1227 TreeJsonTarget::Workspace(_) => {
1228 self.ensure_workspace();
1229 for (package, edge) in edges {
1230 let Edge::Dev(group, ..) = edge else {
1231 continue;
1232 };
1233 let group = self.ensure_workspace_group(group);
1234 let marker = self.marker(edge);
1235 for package in self.dependency_targets(package, edge.extras()) {
1236 self.add_link(group.clone(), package, JsonLink::Dependency(marker.clone()));
1237 }
1238 }
1239 }
1240 }
1241 }
1242
1243 fn marker(&self, edge: &Edge<'_>) -> Option<String> {
1244 self.tree
1245 .lock
1246 .simplify_environment(edge.marker().pep508())
1247 .try_to_string()
1248 }
1249
1250 fn add_link(&mut self, source: String, target: String, link: JsonLink) {
1251 let (source, target) = if self.tree.invert && matches!(&link, JsonLink::Dependency(_)) {
1255 (target, source)
1256 } else {
1257 (source, target)
1258 };
1259 let Some(node) = self.resolution.get_mut(&source) else {
1260 return;
1261 };
1262 match link {
1263 JsonLink::Dependency(marker) => {
1264 node.add_resolution_dependency(target, marker);
1265 }
1266 JsonLink::Optional(name) => {
1267 node.add_optional_dependency(name, target);
1268 }
1269 JsonLink::Group(name) => {
1270 node.add_dependency_group(name, target);
1271 }
1272 }
1273 }
1274
1275 fn add_package_roots(&mut self, roots: &mut Vec<JsonRoot>, package_id: &'env PackageId) {
1276 let package = self.tree.lock.find_by_id(package_id);
1277 let extras = package
1278 .optional_dependencies
1279 .keys()
1280 .cloned()
1281 .collect::<Vec<_>>();
1282 let groups = package
1283 .dependency_groups
1284 .keys()
1285 .filter(|group| self.tree.groups.contains(group))
1286 .cloned()
1287 .collect::<Vec<_>>();
1288
1289 if self.tree.prod {
1290 roots.push(JsonRoot {
1291 id: self.ensure_package(package_id, MetadataNodeKind::Package),
1292 });
1293 for extra in &extras {
1294 let id = MetadataNodeId::from_package_id(
1295 &self.workspace_root,
1296 package_id,
1297 MetadataNodeKind::Extra(extra.clone()),
1298 )
1299 .to_flat();
1300 if self.resolution.contains_key(&id) {
1301 roots.push(JsonRoot { id });
1302 }
1303 }
1304 }
1305
1306 for group in &groups {
1307 let id = MetadataNodeId::from_package_id(
1308 &self.workspace_root,
1309 package_id,
1310 MetadataNodeKind::Group(group.clone()),
1311 )
1312 .to_flat();
1313 if self.resolution.contains_key(&id) {
1314 roots.push(JsonRoot { id });
1315 }
1316 }
1317 }
1318
1319 fn roots(&mut self, target: TreeJsonTarget<'_>) -> Vec<JsonRoot> {
1320 let mut roots = Vec::new();
1321 for root in &self.tree.roots {
1322 match self.tree.graph[*root] {
1323 Node::Package(package_id) => {
1324 if self.tree.invert {
1325 roots.push(JsonRoot {
1326 id: self.ensure_package(package_id, MetadataNodeKind::Package),
1327 });
1328 } else {
1329 self.add_package_roots(&mut roots, package_id);
1330 }
1331 }
1332 Node::Root => match target {
1333 TreeJsonTarget::Script(path) => {
1334 let script = self.ensure_script(path);
1335 roots.push(JsonRoot { id: script });
1336 }
1337 TreeJsonTarget::Workspace(_) => {
1338 let packages = self
1339 .tree
1340 .graph
1341 .edges_directed(*root, Direction::Outgoing)
1342 .filter(|edge| matches!(edge.weight(), Edge::Prod(..)))
1343 .filter_map(|edge| {
1344 let Node::Package(package_id) = self.tree.graph[edge.target()]
1345 else {
1346 return None;
1347 };
1348 Some(package_id)
1349 })
1350 .collect::<Vec<_>>();
1351 for package_id in packages {
1352 self.add_package_roots(&mut roots, package_id);
1353 }
1354 let groups = self
1355 .tree
1356 .lock
1357 .dependency_groups()
1358 .keys()
1359 .filter(|group| self.tree.groups.contains(group))
1360 .cloned()
1361 .collect::<Vec<_>>();
1362 for group in groups {
1363 let id = MetadataNodeId::from_workspace_group(
1364 self.workspace_root.clone(),
1365 group,
1366 )
1367 .to_flat();
1368 if self.resolution.contains_key(&id) {
1369 roots.push(JsonRoot { id });
1370 }
1371 }
1372 }
1373 },
1374 }
1375 }
1376 roots.sort();
1377 roots.dedup();
1378 roots
1379 }
1380
1381 fn members(&self, target: TreeJsonTarget<'_>) -> Vec<MetadataWorkspaceMember> {
1382 if matches!(target, TreeJsonTarget::Script(_)) {
1383 return Vec::new();
1384 }
1385
1386 let packages = if self.tree.lock.members().is_empty() {
1387 self.tree.lock.root().into_iter().collect::<Vec<_>>()
1388 } else {
1389 self.tree
1390 .lock
1391 .packages()
1392 .iter()
1393 .filter(|package| self.tree.lock.members().contains(&package.id.name))
1394 .collect::<Vec<_>>()
1395 };
1396
1397 packages
1398 .into_iter()
1399 .filter(|package| {
1400 let id = MetadataNodeId::from_package_id(
1401 &self.workspace_root,
1402 &package.id,
1403 MetadataNodeKind::Package,
1404 )
1405 .to_flat();
1406 self.resolution.contains_key(&id)
1407 })
1408 .filter_map(|package| {
1409 MetadataWorkspaceMember::from_locked_package(&self.workspace_root, &package.id)
1410 })
1411 .collect()
1412 }
1413
1414 fn finish(mut self) -> BTreeMap<String, MetadataNode> {
1415 for node in self.resolution.values_mut() {
1416 node.normalize_resolution();
1417 }
1418 self.resolution
1419 }
1420}
1421
1422enum JsonLink {
1423 Dependency(Option<String>),
1424 Optional(ExtraName),
1425 Group(GroupName),
1426}
1427
1428#[derive(Debug, Serialize)]
1429struct JsonSchema {
1430 version: JsonSchemaVersion,
1431}
1432
1433#[derive(Debug, Serialize)]
1434#[serde(rename_all = "snake_case")]
1435enum JsonSchemaVersion {
1436 Preview,
1437}
1438
1439#[derive(Debug, Serialize, PartialEq, Eq, PartialOrd, Ord)]
1440struct JsonRoot {
1441 id: String,
1442}
1443
1444#[derive(Debug, Clone, PartialEq, Eq, Hash)]
1445struct VisitedNode<'env> {
1446 package_id: &'env PackageId,
1447 expanded_extras: BTreeSet<&'env ExtraName>,
1448 marker: Option<UniversalMarker>,
1449}
1450
1451#[derive(Debug, Clone, PartialEq, Eq, Ord, PartialOrd)]
1452enum Node<'env> {
1453 Root,
1455 Package(&'env PackageId),
1457}
1458
1459#[derive(Debug, Clone, PartialEq, Eq, Ord, PartialOrd)]
1460enum Edge<'env> {
1461 Prod(Option<RequestedExtras<'env>>, UniversalMarker),
1462 Optional(
1463 &'env ExtraName,
1464 Option<RequestedExtras<'env>>,
1465 UniversalMarker,
1466 ),
1467 Dev(
1468 &'env GroupName,
1469 Option<RequestedExtras<'env>>,
1470 UniversalMarker,
1471 ),
1472}
1473
1474impl<'env> Edge<'env> {
1475 fn extras(&self) -> Option<RequestedExtras<'env>> {
1476 match self {
1477 Self::Prod(extras, _) => *extras,
1478 Self::Optional(_, extras, _) => *extras,
1479 Self::Dev(_, extras, _) => *extras,
1480 }
1481 }
1482
1483 fn required_extra(&self) -> Option<&'env ExtraName> {
1484 match self {
1485 Self::Optional(extra, ..) => Some(extra),
1486 Self::Prod(..) | Self::Dev(..) => None,
1487 }
1488 }
1489
1490 fn marker(&self) -> UniversalMarker {
1491 match self {
1492 Self::Prod(_, marker) | Self::Optional(_, _, marker) | Self::Dev(_, _, marker) => {
1493 *marker
1494 }
1495 }
1496 }
1497
1498 fn is_dev(&self) -> bool {
1499 matches!(self, Self::Dev(..))
1500 }
1501
1502 fn kind(&self) -> EdgeKind<'env> {
1503 match self {
1504 Self::Prod(..) => EdgeKind::Prod,
1505 Self::Optional(extra, ..) => EdgeKind::Optional(extra),
1506 Self::Dev(group, ..) => EdgeKind::Dev(group),
1507 }
1508 }
1509}
1510
1511#[derive(Debug, Copy, Clone)]
1512enum RequestedExtras<'env> {
1513 Dependency(&'env BTreeSet<ExtraName>),
1514 Requirement(&'env [ExtraName]),
1515}
1516
1517impl PartialEq for RequestedExtras<'_> {
1518 fn eq(&self, other: &Self) -> bool {
1519 self.iter().eq(other.iter())
1520 }
1521}
1522
1523impl Eq for RequestedExtras<'_> {}
1524
1525impl PartialOrd for RequestedExtras<'_> {
1526 fn partial_cmp(&self, other: &Self) -> Option<Ordering> {
1527 Some(self.cmp(other))
1528 }
1529}
1530
1531impl Ord for RequestedExtras<'_> {
1532 fn cmp(&self, other: &Self) -> Ordering {
1533 self.iter().cmp(other.iter())
1534 }
1535}
1536
1537impl<'env> RequestedExtras<'env> {
1538 fn contains(self, extra: &ExtraName) -> bool {
1539 match self {
1540 Self::Dependency(extras) => extras.contains(extra),
1541 Self::Requirement(extras) => extras.contains(extra),
1542 }
1543 }
1544
1545 fn is_empty(self) -> bool {
1546 match self {
1547 Self::Dependency(extras) => extras.is_empty(),
1548 Self::Requirement(extras) => extras.is_empty(),
1549 }
1550 }
1551
1552 fn iter(self) -> impl Iterator<Item = &'env ExtraName> {
1553 match self {
1554 Self::Dependency(extras) => Either::Left(extras.iter()),
1555 Self::Requirement(extras) => Either::Right(extras.iter()),
1556 }
1557 }
1558}
1559
1560#[derive(Debug, Clone, PartialEq, Eq, Ord, PartialOrd)]
1561enum EdgeKind<'env> {
1562 Prod,
1563 Optional(&'env ExtraName),
1564 Dev(&'env GroupName),
1565}
1566
1567#[derive(Debug, Copy, Clone, PartialEq, Eq, Ord, PartialOrd)]
1569struct Cursor(NodeIndex, Option<EdgeIndex>, UniversalMarker);
1570
1571impl Cursor {
1572 fn new(node: NodeIndex, edge: EdgeIndex, marker: UniversalMarker) -> Self {
1574 Self(node, Some(edge), marker)
1575 }
1576
1577 fn root(node: NodeIndex, marker: UniversalMarker) -> Self {
1579 Self(node, None, marker)
1580 }
1581
1582 fn node(&self) -> NodeIndex {
1584 self.0
1585 }
1586
1587 fn edge(&self) -> Option<EdgeIndex> {
1589 self.1
1590 }
1591
1592 fn marker(&self) -> UniversalMarker {
1594 self.2
1595 }
1596}
1597
1598impl std::fmt::Display for TreeDisplay<'_> {
1599 fn fmt(&self, f: &mut std::fmt::Formatter) -> std::fmt::Result {
1600 use owo_colors::OwoColorize;
1601
1602 let mut deduped = false;
1603 for line in self.render() {
1604 deduped |= line.contains('*');
1605 writeln!(f, "{line}")?;
1606 }
1607
1608 if deduped {
1609 let message = if self.no_dedupe {
1610 "(*) Package tree is a cycle and cannot be shown".italic()
1611 } else {
1612 "(*) Package tree already displayed".italic()
1613 };
1614 writeln!(f, "{message}")?;
1615 }
1616
1617 Ok(())
1618 }
1619}