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, PackageIndex};
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, 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 = vec![None; size_guess];
119 let mut queue: VecDeque<(PackageIndex, 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 package_index = lock.by_id[id];
131 let dist = lock.package(package_index);
132
133 let index = *inverse[package_index.0]
137 .get_or_insert_with(|| graph.add_node(Node::Package(package_index)));
138
139 graph.add_edge(root, index, Edge::Prod(None, UniversalMarker::TRUE));
141
142 if groups.prod() {
143 if seen.insert((package_index, None)) {
145 queue.push_back((package_index, None));
146 }
147
148 for extra in dist.optional_dependencies.keys() {
150 if seen.insert((package_index, Some(extra))) {
151 queue.push_back((package_index, 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[dep.index.0]
181 .get_or_insert_with(|| graph.add_node(Node::Package(dep.index)));
182
183 graph.add_edge(
185 index,
186 dep_index,
187 Edge::Dev(
188 group,
189 Some(RequestedExtras::Dependency(&dep.extra)),
190 dep.complexified_marker,
191 ),
192 );
193
194 if seen.insert((dep.index, None)) {
196 queue.push_back((dep.index, None));
197 }
198 for extra in &dep.extra {
199 if seen.insert((dep.index, Some(extra))) {
200 queue.push_back((dep.index, Some(extra)));
201 }
202 }
203 }
204 }
205
206 {
215 let by_name: FxHashMap<_, Vec<_>> = {
217 lock.packages().iter().enumerate().fold(
218 FxHashMap::with_capacity_and_hasher(lock.len(), FxBuildHasher),
219 |mut map, (index, package)| {
220 map.entry(&package.id.name)
221 .or_default()
222 .push(PackageIndex(index));
223 map
224 },
225 )
226 };
227
228 for requirement in lock.requirements() {
230 for &package_index in by_name.get(&requirement.name).into_iter().flatten() {
231 let package = lock.package(package_index);
232 let marker = if package.fork_markers.is_empty() {
235 requirement.marker
236 } else {
237 let mut combined = MarkerTree::FALSE;
238 for fork_marker in &package.fork_markers {
239 combined = combined.or(fork_marker.pep508());
240 }
241 combined = combined.and(requirement.marker);
242 combined
243 };
244 if marker.is_false() {
245 continue;
246 }
247 if markers.is_some_and(|markers| !marker.evaluate(markers, &[])) {
248 continue;
249 }
250 let index = *inverse[package_index.0]
252 .get_or_insert_with(|| graph.add_node(Node::Package(package_index)));
253
254 graph.add_edge(
256 root,
257 index,
258 Edge::Prod(
259 Some(RequestedExtras::Requirement(requirement.extras.as_ref())),
260 UniversalMarker::from_combined(marker),
261 ),
262 );
263
264 if seen.insert((package_index, None)) {
266 queue.push_back((package_index, None));
267 }
268 for extra in &*requirement.extras {
269 if seen.insert((package_index, Some(extra))) {
270 queue.push_back((package_index, Some(extra)));
271 }
272 }
273 }
274 }
275
276 for (group, requirements) in lock.dependency_groups() {
278 if !groups.contains(group) {
279 continue;
280 }
281 for requirement in requirements {
282 for &package_index in by_name.get(&requirement.name).into_iter().flatten() {
283 let package = lock.package(package_index);
284 let marker = if package.fork_markers.is_empty() {
287 requirement.marker
288 } else {
289 let mut combined = MarkerTree::FALSE;
290 for fork_marker in &package.fork_markers {
291 combined = combined.or(fork_marker.pep508());
292 }
293 combined = combined.and(requirement.marker);
294 combined
295 };
296 if marker.is_false() {
297 continue;
298 }
299 if markers.is_some_and(|markers| !marker.evaluate(markers, &[])) {
300 continue;
301 }
302 let index = *inverse[package_index.0]
304 .get_or_insert_with(|| graph.add_node(Node::Package(package_index)));
305
306 graph.add_edge(
308 root,
309 index,
310 Edge::Dev(
311 group,
312 Some(RequestedExtras::Requirement(requirement.extras.as_ref())),
313 UniversalMarker::from_combined(marker),
314 ),
315 );
316
317 if seen.insert((package_index, None)) {
319 queue.push_back((package_index, None));
320 }
321 for extra in &*requirement.extras {
322 if seen.insert((package_index, Some(extra))) {
323 queue.push_back((package_index, Some(extra)));
324 }
325 }
326 }
327 }
328 }
329 }
330
331 while let Some((package_index, extra)) = queue.pop_front() {
333 let index = inverse[package_index.0].expect("queued package has a graph node");
334 let package = lock.package(package_index);
335
336 let deps = if let Some(extra) = extra {
337 Either::Left(
338 package
339 .optional_dependencies
340 .get(extra)
341 .into_iter()
342 .flatten(),
343 )
344 } else {
345 Either::Right(package.dependencies.iter())
346 };
347
348 for dep in deps {
349 if prune.contains(&dep.package_id.name) {
350 continue;
351 }
352
353 if markers
354 .is_some_and(|markers| !dep.complexified_marker.evaluate_no_extras(markers))
355 {
356 continue;
357 }
358
359 let dep_index = *inverse[dep.index.0]
361 .get_or_insert_with(|| graph.add_node(Node::Package(dep.index)));
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.index, None)) {
383 queue.push_back((dep.index, None));
384 }
385 for extra in &dep.extra {
386 if seen.insert((dep.index, Some(extra))) {
387 queue.push_back((dep.index, 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_index) => {
399 members.contains(&lock.package(package_index).id)
400 }
401 Node::Root => true,
402 })
403 .collect::<FxHashSet<_>>();
404 let mut stack = reachable.iter().copied().collect::<VecDeque<_>>();
405 while let Some(node) = stack.pop_front() {
406 for edge in graph.edges_directed(node, Direction::Outgoing) {
407 if reachable.insert(edge.target()) {
408 stack.push_back(edge.target());
409 }
410 }
411 }
412
413 graph.retain_nodes(|_, index| reachable.contains(&index));
415 }
416
417 if invert {
419 graph.reverse();
420 }
421
422 if !packages.is_empty() {
424 let mut reachable = graph
425 .node_indices()
426 .filter(|index| {
427 let Node::Package(package_index) = graph[*index] else {
428 return false;
429 };
430 packages.contains(&lock.package(package_index).id.name)
431 })
432 .collect::<FxHashSet<_>>();
433 let mut stack = reachable.iter().copied().collect::<VecDeque<_>>();
434 while let Some(node) = stack.pop_front() {
435 for edge in graph.edges_directed(node, Direction::Outgoing) {
436 if reachable.insert(edge.target()) {
437 stack.push_back(edge.target());
438 }
439 }
440 }
441
442 graph.retain_nodes(|_, index| reachable.contains(&index));
444 }
445
446 let roots = {
448 if !packages.is_empty() {
450 let mut roots = graph
451 .node_indices()
452 .filter(|index| {
453 let Node::Package(package_index) = graph[*index] else {
454 return false;
455 };
456 packages.contains(&lock.package(package_index).id.name)
457 })
458 .collect::<Vec<_>>();
459
460 roots.sort_by_key(|index| graph[*index].sort_key(lock));
462
463 roots
464 } else {
465 let mut roots = if invert {
466 graph
469 .node_indices()
470 .filter(|index| {
471 graph
472 .edges_directed(*index, Direction::Incoming)
473 .next()
474 .is_none()
475 })
476 .collect::<Vec<_>>()
477 } else {
478 graph
480 .node_indices()
481 .filter(|index| matches!(graph[*index], Node::Root))
482 .collect::<Vec<_>>()
483 };
484
485 roots.sort_by_key(|index| graph[*index].sort_key(lock));
486 roots
487 }
488 };
489
490 Self {
491 graph,
492 roots,
493 latest,
494 depth,
495 no_dedupe,
496 invert,
497 prod: groups.prod(),
498 groups: groups.clone(),
499 lock,
500 show_sizes,
501 conflict_marker,
502 }
503 }
504
505 fn visit(
507 &'env self,
508 cursor: Cursor,
509 visited: &mut FxHashMap<VisitedNode<'env>, Vec<PackageIndex>>,
510 path: &mut Vec<VisitedNode<'env>>,
511 ) -> Vec<String> {
512 if path.len() > self.depth {
514 return Vec::new();
515 }
516
517 let Node::Package(package_index) = self.graph[cursor.node()] else {
518 return Vec::new();
519 };
520 let edge = cursor.edge().map(|edge_id| &self.graph[edge_id]);
521 let package = self.lock.package(package_index);
522 let package_id = &package.id;
523
524 let expanded_extras = self.expanded_extras(package, edge);
525 let visited_node = VisitedNode {
526 package_index,
527 expanded_extras: expanded_extras.clone(),
528 marker: self.invert.then_some(cursor.marker()),
529 };
530
531 let line = {
532 let mut line = format!("{}", package_id.name);
533
534 if let Some(extras) = edge.and_then(Edge::extras) {
535 if !extras.is_empty() {
536 line.push('[');
537 line.push_str(extras.iter().join(", ").as_str());
538 line.push(']');
539 }
540 }
541
542 if let Some(version) = package_id.version.as_ref() {
543 line.push(' ');
544 line.push('v');
545 let _ = write!(line, "{version}");
546 }
547
548 if let Some(edge) = edge {
549 match edge {
550 Edge::Prod(..) => {}
551 Edge::Optional(extra, ..) => {
552 let _ = write!(line, " (extra: {extra})");
553 }
554 Edge::Dev(group, ..) => {
555 let _ = write!(line, " (group: {group})");
556 }
557 }
558 }
559
560 if self.show_sizes {
563 if let Some(size_bytes) = package.wheels.iter().find_map(|wheel| wheel.size) {
564 let bytes = human_readable_bytes(size_bytes);
565 line.push(' ');
566 line.push_str(format!("{}", format!("({bytes:.1})").dimmed()).as_str());
567 }
568 }
569
570 line
571 };
572
573 if path.contains(&visited_node) {
577 return vec![format!("{line} (*)")];
578 }
579 if !self.no_dedupe
580 && let Some(requirements) = visited.get(&visited_node)
581 {
582 return if requirements.is_empty() {
583 vec![line]
584 } else {
585 vec![format!("{line} (*)")]
586 };
587 }
588
589 let line = if let Some(version) = self.latest.get(package_id) {
591 format!("{line} {}", format!("(latest: v{version})").bold().cyan())
592 } else {
593 line
594 };
595
596 let mut dependencies = if self.invert && edge.is_some_and(Edge::is_dev) {
597 Vec::new()
600 } else {
601 self.graph
602 .edges_directed(cursor.node(), Direction::Outgoing)
603 .filter_map(|edge| match self.graph[edge.target()] {
604 Node::Root => None,
605 Node::Package(_) => {
606 let edge_kind = &self.graph[edge.id()];
607
608 if self.invert {
609 if !expanded_extras.is_empty()
612 && edge_kind.extras().is_none_or(|extras| {
613 !expanded_extras.iter().all(|extra| extras.contains(extra))
614 })
615 {
616 return None;
617 }
618
619 let mut marker = cursor.marker();
622 marker.and(edge_kind.marker());
623 if marker.is_false() {
624 return None;
625 }
626 Some(Cursor::new(edge.target(), edge.id(), marker))
627 } else {
628 if let Edge::Optional(required_extra, ..) = edge_kind
631 && !expanded_extras.contains(required_extra)
632 {
633 return None;
634 }
635 Some(Cursor::new(edge.target(), edge.id(), UniversalMarker::TRUE))
636 }
637 }
638 })
639 .collect::<Vec<_>>()
640 };
641 dependencies.sort_by_key(|cursor| {
642 let node = self.graph[cursor.node()].sort_key(self.lock);
643 let edge = cursor
644 .edge()
645 .map(|edge_id| &self.graph[edge_id])
646 .map(Edge::kind);
647 (edge, node)
648 });
649
650 let mut lines = vec![line];
651
652 if path.len() < self.depth {
655 visited.insert(
656 visited_node.clone(),
657 dependencies
658 .iter()
659 .filter_map(|node| match self.graph[node.node()] {
660 Node::Package(package_index) => Some(package_index),
661 Node::Root => None,
662 })
663 .collect(),
664 );
665 }
666 path.push(visited_node);
667
668 for (index, dep) in dependencies.iter().enumerate() {
669 let (prefix_top, prefix_rest) = if dependencies.len() - 1 == index {
689 ("└── ", " ")
690 } else {
691 ("├── ", "│ ")
692 };
693 for (visited_index, visited_line) in self.visit(*dep, visited, path).iter().enumerate()
694 {
695 let prefix = if visited_index == 0 {
696 prefix_top
697 } else {
698 prefix_rest
699 };
700 lines.push(format!("{prefix}{visited_line}"));
701 }
702 }
703
704 path.pop();
705
706 lines
707 }
708
709 fn render(&self) -> Vec<String> {
711 let mut path = Vec::new();
712 let mut lines = Vec::with_capacity(self.graph.node_count());
713 let mut visited =
714 FxHashMap::with_capacity_and_hasher(self.graph.node_count(), FxBuildHasher);
715
716 for node in &self.roots {
717 match self.graph[*node] {
718 Node::Root => {
719 for edge in self.graph.edges_directed(*node, Direction::Outgoing) {
720 let node = edge.target();
721 path.clear();
722 lines.extend(self.visit(
723 Cursor::new(node, edge.id(), self.conflict_marker),
724 &mut visited,
725 &mut path,
726 ));
727 }
728 }
729 Node::Package(_) => {
730 path.clear();
731 lines.extend(self.visit(
732 Cursor::root(*node, self.conflict_marker),
733 &mut visited,
734 &mut path,
735 ));
736 }
737 }
738 }
739
740 lines
741 }
742
743 fn expanded_extras(
745 &self,
746 package: &'env Package,
747 edge: Option<&Edge<'env>>,
748 ) -> BTreeSet<&'env ExtraName> {
749 if self.invert {
750 return edge.and_then(Edge::required_extra).into_iter().collect();
753 }
754
755 let Some(requested_extras) = edge.and_then(Edge::extras) else {
756 return package.optional_dependencies.keys().collect();
758 };
759
760 requested_extras
761 .iter()
762 .filter(|extra| package.optional_dependencies.contains_key(*extra))
763 .collect()
764 }
765
766 pub fn to_json(&self, target: TreeJsonTarget<'_>) -> Result<String, serde_json::Error> {
768 serde_json::to_string_pretty(&JsonGraph::new(self, target))
769 }
770
771 fn json_traversal(&self) -> JsonTraversal {
781 let mut distances = FxHashMap::default();
782 let mut queue = VecDeque::new();
783 let mut nodes = FxHashSet::default();
784 let mut edges = FxHashSet::default();
785
786 for root in &self.roots {
787 match self.graph[*root] {
788 Node::Root => {
789 for edge in self.graph.edges_directed(*root, Direction::Outgoing) {
790 let Node::Package(package_index) = self.graph[edge.target()] else {
791 continue;
792 };
793 let state = JsonTraversalNode {
794 index: edge.target(),
795 expanded_extras: self.expanded_extras(
796 self.lock.package(package_index),
797 Some(edge.weight()),
798 ),
799 marker: UniversalMarker::TRUE,
800 reached_via_dependency_group: false,
801 };
802 nodes.insert(state.index);
803 if distances.insert(state.clone(), 0).is_none() {
804 queue.push_back(state);
805 }
806 }
807 }
808 Node::Package(package_index) => {
809 let state = JsonTraversalNode {
810 index: *root,
811 expanded_extras: self
812 .expanded_extras(self.lock.package(package_index), None),
813 marker: if self.invert {
814 self.conflict_marker
815 } else {
816 UniversalMarker::TRUE
817 },
818 reached_via_dependency_group: false,
819 };
820 nodes.insert(state.index);
821 if distances.insert(state.clone(), 0).is_none() {
822 queue.push_back(state);
823 }
824 }
825 }
826 }
827
828 while let Some(source) = queue.pop_front() {
829 let distance = distances[&source];
830 if distance >= self.depth || self.invert && source.reached_via_dependency_group {
831 continue;
832 }
833
834 for edge in self.graph.edges_directed(source.index, Direction::Outgoing) {
835 let edge_kind = edge.weight();
836 let marker = if self.invert {
837 if !source.expanded_extras.is_empty()
840 && edge_kind.extras().is_none_or(|extras| {
841 !source
842 .expanded_extras
843 .iter()
844 .all(|extra| extras.contains(extra))
845 })
846 {
847 continue;
848 }
849
850 let mut marker = source.marker;
853 marker.and(edge_kind.marker());
854 if marker.is_false() {
855 continue;
856 }
857 marker
858 } else {
859 if let Edge::Optional(required_extra, ..) = edge_kind
862 && !source.expanded_extras.contains(required_extra)
863 {
864 continue;
865 }
866 UniversalMarker::TRUE
867 };
868
869 let target = edge.target();
870 if matches!(self.graph[target], Node::Root) {
871 edges.insert(edge.id());
872 continue;
873 }
874 let Node::Package(package_index) = self.graph[target] else {
875 continue;
876 };
877 let state = JsonTraversalNode {
878 index: target,
879 expanded_extras: self
880 .expanded_extras(self.lock.package(package_index), Some(edge.weight())),
881 marker,
882 reached_via_dependency_group: self.invert && edge_kind.is_dev(),
883 };
884 nodes.insert(state.index);
885 edges.insert(edge.id());
886 if !distances.contains_key(&state) {
887 distances.insert(state.clone(), distance + 1);
888 queue.push_back(state);
889 }
890 }
891 }
892
893 JsonTraversal { nodes, edges }
894 }
895}
896
897#[derive(Debug)]
898struct JsonTraversal {
899 nodes: FxHashSet<NodeIndex>,
900 edges: FxHashSet<EdgeIndex>,
901}
902
903#[derive(Debug, Clone, PartialEq, Eq, Hash)]
904struct JsonTraversalNode<'env> {
905 index: NodeIndex,
906 expanded_extras: BTreeSet<&'env ExtraName>,
907 marker: UniversalMarker,
908 reached_via_dependency_group: bool,
909}
910
911#[derive(Debug, Serialize)]
1011struct JsonGraph {
1012 schema: JsonSchema,
1013 workspace_root: PortablePathBuf,
1014 #[serde(skip_serializing_if = "Option::is_none")]
1015 script: Option<MetadataScript>,
1016 #[serde(skip_serializing_if = "Option::is_none")]
1017 workspace: Option<MetadataWorkspace>,
1018 roots: Vec<JsonRoot>,
1019 inverted: bool,
1020 #[serde(skip_serializing_if = "Vec::is_empty")]
1022 members: Vec<MetadataWorkspaceMember>,
1023 resolution: BTreeMap<String, MetadataNode>,
1024}
1025
1026impl JsonGraph {
1027 fn new(tree: &TreeDisplay<'_>, target: TreeJsonTarget<'_>) -> Self {
1028 let traversal = tree.json_traversal();
1029 let workspace_root = PortablePathBuf::from(target.root());
1030 let mut builder = JsonGraphBuilder::new(tree, workspace_root.clone());
1031
1032 for index in traversal.nodes.iter().copied() {
1033 let Node::Package(package_index) = tree.graph[index] else {
1034 continue;
1035 };
1036 builder.ensure_package(package_index, MetadataNodeKind::Package);
1037 }
1038
1039 for edge in tree
1040 .graph
1041 .edge_references()
1042 .filter(|edge| traversal.edges.contains(&edge.id()))
1043 {
1044 builder.add_package_edge(edge.source(), edge.target(), edge.weight());
1045 }
1046
1047 builder.add_target_edges(target, &traversal);
1048 let (script, workspace) = match target {
1049 TreeJsonTarget::Script(path) => {
1050 let path = PortablePathBuf::from(path);
1051 let id = builder.ensure_script(path.as_ref());
1052 (Some(MetadataScript::new(path, id)), None)
1053 }
1054 TreeJsonTarget::Workspace(path) => {
1055 let path = PortablePathBuf::from(path);
1056 let id = builder.ensure_workspace();
1057 (None, Some(MetadataWorkspace::new(path, id)))
1058 }
1059 };
1060 let roots = builder.roots(target);
1061 let members = builder.members(target);
1062 let resolution = builder.finish();
1063
1064 Self {
1065 schema: JsonSchema {
1066 version: JsonSchemaVersion::Preview,
1067 },
1068 workspace_root,
1069 script,
1070 workspace,
1071 roots,
1072 inverted: tree.invert,
1073 members,
1074 resolution,
1075 }
1076 }
1077}
1078
1079struct JsonGraphBuilder<'tree, 'env> {
1080 tree: &'tree TreeDisplay<'env>,
1081 workspace_root: PortablePathBuf,
1082 resolution: BTreeMap<String, MetadataNode>,
1083}
1084
1085impl<'tree, 'env> JsonGraphBuilder<'tree, 'env> {
1086 fn new(tree: &'tree TreeDisplay<'env>, workspace_root: PortablePathBuf) -> Self {
1087 Self {
1088 tree,
1089 workspace_root,
1090 resolution: BTreeMap::new(),
1091 }
1092 }
1093
1094 fn ensure_node(&mut self, identity: MetadataNodeId) -> String {
1095 let id = identity.to_flat();
1096 self.resolution
1097 .entry(id.clone())
1098 .or_insert_with(|| MetadataNode::new(identity));
1099 id
1100 }
1101
1102 fn ensure_package(&mut self, package_index: PackageIndex, kind: MetadataNodeKind) -> String {
1103 let package = self.tree.lock.package(package_index);
1104 let package_id = &package.id;
1105 let is_package = matches!(kind, MetadataNodeKind::Package);
1106 let id = MetadataNodeId::from_package_id(&self.workspace_root, package_id, kind.clone())
1107 .to_flat();
1108 self.resolution.entry(id.clone()).or_insert_with(|| {
1109 let mut node = MetadataNode::from_package_id(&self.workspace_root, package_id, kind);
1110 if is_package {
1111 node.set_latest_version(self.tree.latest.get(package_id).cloned());
1112 node.set_wheels_from_package(&self.workspace_root, package);
1113 }
1114 node
1115 });
1116 id
1117 }
1118
1119 fn ensure_extra(&mut self, package_index: PackageIndex, extra: &ExtraName) -> String {
1120 let package = self.ensure_package(package_index, MetadataNodeKind::Package);
1121 let extra_id = self.ensure_package(package_index, MetadataNodeKind::Extra(extra.clone()));
1122 self.add_link(
1123 package.clone(),
1124 extra_id.clone(),
1125 JsonLink::Optional(extra.clone()),
1126 );
1127 self.add_link(extra_id.clone(), package, JsonLink::Dependency(None));
1128 extra_id
1129 }
1130
1131 fn ensure_group(&mut self, package_index: PackageIndex, group: &GroupName) -> String {
1132 let package = self.ensure_package(package_index, MetadataNodeKind::Package);
1133 let group_id = self.ensure_package(package_index, MetadataNodeKind::Group(group.clone()));
1134 self.add_link(package, group_id.clone(), JsonLink::Group(group.clone()));
1135 group_id
1136 }
1137
1138 fn ensure_workspace(&mut self) -> String {
1139 self.ensure_node(MetadataNodeId::from_workspace(self.workspace_root.clone()))
1140 }
1141
1142 fn ensure_workspace_group(&mut self, group: &GroupName) -> String {
1143 let workspace = self.ensure_workspace();
1144 let group_id = self.ensure_node(MetadataNodeId::from_workspace_group(
1145 self.workspace_root.clone(),
1146 group.clone(),
1147 ));
1148 self.add_link(workspace, group_id.clone(), JsonLink::Group(group.clone()));
1149 group_id
1150 }
1151
1152 fn ensure_script(&mut self, path: &Path) -> String {
1153 self.ensure_node(MetadataNodeId::from_script(PortablePathBuf::from(path)))
1154 }
1155
1156 fn dependency_targets(
1157 &mut self,
1158 package_index: PackageIndex,
1159 extras: Option<RequestedExtras<'env>>,
1160 ) -> Vec<String> {
1161 let Some(extras) = extras.filter(|extras| !extras.is_empty()) else {
1162 return vec![self.ensure_package(package_index, MetadataNodeKind::Package)];
1163 };
1164 extras
1165 .iter()
1166 .map(|extra| self.ensure_extra(package_index, extra))
1167 .collect()
1168 }
1169
1170 fn add_package_edge(&mut self, source: NodeIndex, target: NodeIndex, edge: &Edge<'env>) {
1171 let (source, target) = if self.tree.invert {
1172 (target, source)
1173 } else {
1174 (source, target)
1175 };
1176 let (Node::Package(source), Node::Package(target)) =
1177 (&self.tree.graph[source], &self.tree.graph[target])
1178 else {
1179 return;
1180 };
1181 let (source, target) = (*source, *target);
1182
1183 let source = match edge {
1184 Edge::Prod(..) => self.ensure_package(source, MetadataNodeKind::Package),
1185 Edge::Optional(extra, ..) => self.ensure_extra(source, extra),
1186 Edge::Dev(group, ..) => self.ensure_group(source, group),
1187 };
1188 let marker = self.marker(edge);
1189 for target in self.dependency_targets(target, edge.extras()) {
1190 self.add_link(source.clone(), target, JsonLink::Dependency(marker.clone()));
1191 }
1192 }
1193
1194 fn add_target_edges(&mut self, target: TreeJsonTarget<'_>, traversal: &JsonTraversal) {
1195 let edges = self
1199 .tree
1200 .graph
1201 .edge_references()
1202 .filter(|edge| !self.tree.invert || traversal.edges.contains(&edge.id()))
1203 .filter_map(|edge| {
1204 let package = match (
1205 &self.tree.graph[edge.source()],
1206 &self.tree.graph[edge.target()],
1207 ) {
1208 (Node::Root, Node::Package(package)) | (Node::Package(package), Node::Root) => {
1209 *package
1210 }
1211 (Node::Root, Node::Root) | (Node::Package(_), Node::Package(_)) => return None,
1212 };
1213 Some((package, edge.weight()))
1214 })
1215 .collect::<Vec<_>>();
1216
1217 match target {
1218 TreeJsonTarget::Script(path) => {
1219 let script = self.ensure_script(path);
1220 for (package, edge) in edges {
1221 let marker = self.marker(edge);
1222 for package in self.dependency_targets(package, edge.extras()) {
1223 self.add_link(
1224 script.clone(),
1225 package,
1226 JsonLink::Dependency(marker.clone()),
1227 );
1228 }
1229 }
1230 }
1231 TreeJsonTarget::Workspace(_) => {
1232 self.ensure_workspace();
1233 for (package, edge) in edges {
1234 let Edge::Dev(group, ..) = edge else {
1235 continue;
1236 };
1237 let group = self.ensure_workspace_group(group);
1238 let marker = self.marker(edge);
1239 for package in self.dependency_targets(package, edge.extras()) {
1240 self.add_link(group.clone(), package, JsonLink::Dependency(marker.clone()));
1241 }
1242 }
1243 }
1244 }
1245 }
1246
1247 fn marker(&self, edge: &Edge<'_>) -> Option<String> {
1248 self.tree
1249 .lock
1250 .simplify_environment(edge.marker().pep508())
1251 .try_to_string()
1252 }
1253
1254 fn add_link(&mut self, source: String, target: String, link: JsonLink) {
1255 let (source, target) = if self.tree.invert && matches!(&link, JsonLink::Dependency(_)) {
1259 (target, source)
1260 } else {
1261 (source, target)
1262 };
1263 let Some(node) = self.resolution.get_mut(&source) else {
1264 return;
1265 };
1266 match link {
1267 JsonLink::Dependency(marker) => {
1268 node.add_resolution_dependency(target, marker);
1269 }
1270 JsonLink::Optional(name) => {
1271 node.add_optional_dependency(name, target);
1272 }
1273 JsonLink::Group(name) => {
1274 node.add_dependency_group(name, target);
1275 }
1276 }
1277 }
1278
1279 fn add_package_roots(&mut self, roots: &mut Vec<JsonRoot>, package_index: PackageIndex) {
1280 let package = self.tree.lock.package(package_index);
1281 let package_id = &package.id;
1282 let extras = package
1283 .optional_dependencies
1284 .keys()
1285 .cloned()
1286 .collect::<Vec<_>>();
1287 let groups = package
1288 .dependency_groups
1289 .keys()
1290 .filter(|group| self.tree.groups.contains(group))
1291 .cloned()
1292 .collect::<Vec<_>>();
1293
1294 if self.tree.prod {
1295 roots.push(JsonRoot {
1296 id: self.ensure_package(package_index, MetadataNodeKind::Package),
1297 });
1298 for extra in &extras {
1299 let id = MetadataNodeId::from_package_id(
1300 &self.workspace_root,
1301 package_id,
1302 MetadataNodeKind::Extra(extra.clone()),
1303 )
1304 .to_flat();
1305 if self.resolution.contains_key(&id) {
1306 roots.push(JsonRoot { id });
1307 }
1308 }
1309 }
1310
1311 for group in &groups {
1312 let id = MetadataNodeId::from_package_id(
1313 &self.workspace_root,
1314 package_id,
1315 MetadataNodeKind::Group(group.clone()),
1316 )
1317 .to_flat();
1318 if self.resolution.contains_key(&id) {
1319 roots.push(JsonRoot { id });
1320 }
1321 }
1322 }
1323
1324 fn roots(&mut self, target: TreeJsonTarget<'_>) -> Vec<JsonRoot> {
1325 let mut roots = Vec::new();
1326 for root in &self.tree.roots {
1327 match self.tree.graph[*root] {
1328 Node::Package(package_index) => {
1329 if self.tree.invert {
1330 roots.push(JsonRoot {
1331 id: self.ensure_package(package_index, MetadataNodeKind::Package),
1332 });
1333 } else {
1334 self.add_package_roots(&mut roots, package_index);
1335 }
1336 }
1337 Node::Root => match target {
1338 TreeJsonTarget::Script(path) => {
1339 let script = self.ensure_script(path);
1340 roots.push(JsonRoot { id: script });
1341 }
1342 TreeJsonTarget::Workspace(_) => {
1343 let packages = self
1344 .tree
1345 .graph
1346 .edges_directed(*root, Direction::Outgoing)
1347 .filter(|edge| matches!(edge.weight(), Edge::Prod(..)))
1348 .filter_map(|edge| {
1349 let Node::Package(package_index) = self.tree.graph[edge.target()]
1350 else {
1351 return None;
1352 };
1353 Some(package_index)
1354 })
1355 .collect::<Vec<_>>();
1356 for package_index in packages {
1357 self.add_package_roots(&mut roots, package_index);
1358 }
1359 let groups = self
1360 .tree
1361 .lock
1362 .dependency_groups()
1363 .keys()
1364 .filter(|group| self.tree.groups.contains(group))
1365 .cloned()
1366 .collect::<Vec<_>>();
1367 for group in groups {
1368 let id = MetadataNodeId::from_workspace_group(
1369 self.workspace_root.clone(),
1370 group,
1371 )
1372 .to_flat();
1373 if self.resolution.contains_key(&id) {
1374 roots.push(JsonRoot { id });
1375 }
1376 }
1377 }
1378 },
1379 }
1380 }
1381 roots.sort();
1382 roots.dedup();
1383 roots
1384 }
1385
1386 fn members(&self, target: TreeJsonTarget<'_>) -> Vec<MetadataWorkspaceMember> {
1387 if matches!(target, TreeJsonTarget::Script(_)) {
1388 return Vec::new();
1389 }
1390
1391 let packages = if self.tree.lock.members().is_empty() {
1392 self.tree.lock.root().into_iter().collect::<Vec<_>>()
1393 } else {
1394 self.tree
1395 .lock
1396 .packages()
1397 .iter()
1398 .filter(|package| self.tree.lock.members().contains(&package.id.name))
1399 .collect::<Vec<_>>()
1400 };
1401
1402 packages
1403 .into_iter()
1404 .filter(|package| {
1405 let id = MetadataNodeId::from_package_id(
1406 &self.workspace_root,
1407 &package.id,
1408 MetadataNodeKind::Package,
1409 )
1410 .to_flat();
1411 self.resolution.contains_key(&id)
1412 })
1413 .filter_map(|package| {
1414 MetadataWorkspaceMember::from_locked_package(&self.workspace_root, &package.id)
1415 })
1416 .collect()
1417 }
1418
1419 fn finish(mut self) -> BTreeMap<String, MetadataNode> {
1420 for node in self.resolution.values_mut() {
1421 node.normalize_resolution();
1422 }
1423 self.resolution
1424 }
1425}
1426
1427enum JsonLink {
1428 Dependency(Option<String>),
1429 Optional(ExtraName),
1430 Group(GroupName),
1431}
1432
1433#[derive(Debug, Serialize)]
1434struct JsonSchema {
1435 version: JsonSchemaVersion,
1436}
1437
1438#[derive(Debug, Serialize)]
1439#[serde(rename_all = "snake_case")]
1440enum JsonSchemaVersion {
1441 Preview,
1442}
1443
1444#[derive(Debug, Serialize, PartialEq, Eq, PartialOrd, Ord)]
1445struct JsonRoot {
1446 id: String,
1447}
1448
1449#[derive(Debug, Clone, PartialEq, Eq, Hash)]
1450struct VisitedNode<'env> {
1451 package_index: PackageIndex,
1452 expanded_extras: BTreeSet<&'env ExtraName>,
1453 marker: Option<UniversalMarker>,
1454}
1455
1456#[derive(Debug, Clone, Copy, PartialEq, Eq)]
1457enum Node {
1458 Root,
1460 Package(PackageIndex),
1462}
1463
1464impl Node {
1465 fn sort_key(self, lock: &Lock) -> Option<&PackageId> {
1467 match self {
1468 Self::Root => None,
1469 Self::Package(index) => Some(&lock.package(index).id),
1470 }
1471 }
1472}
1473
1474#[derive(Debug, Clone, PartialEq, Eq, Ord, PartialOrd)]
1475enum Edge<'env> {
1476 Prod(Option<RequestedExtras<'env>>, UniversalMarker),
1477 Optional(
1478 &'env ExtraName,
1479 Option<RequestedExtras<'env>>,
1480 UniversalMarker,
1481 ),
1482 Dev(
1483 &'env GroupName,
1484 Option<RequestedExtras<'env>>,
1485 UniversalMarker,
1486 ),
1487}
1488
1489impl<'env> Edge<'env> {
1490 fn extras(&self) -> Option<RequestedExtras<'env>> {
1491 match self {
1492 Self::Prod(extras, _) => *extras,
1493 Self::Optional(_, extras, _) => *extras,
1494 Self::Dev(_, extras, _) => *extras,
1495 }
1496 }
1497
1498 fn required_extra(&self) -> Option<&'env ExtraName> {
1499 match self {
1500 Self::Optional(extra, ..) => Some(extra),
1501 Self::Prod(..) | Self::Dev(..) => None,
1502 }
1503 }
1504
1505 fn marker(&self) -> UniversalMarker {
1506 match self {
1507 Self::Prod(_, marker) | Self::Optional(_, _, marker) | Self::Dev(_, _, marker) => {
1508 *marker
1509 }
1510 }
1511 }
1512
1513 fn is_dev(&self) -> bool {
1514 matches!(self, Self::Dev(..))
1515 }
1516
1517 fn kind(&self) -> EdgeKind<'env> {
1518 match self {
1519 Self::Prod(..) => EdgeKind::Prod,
1520 Self::Optional(extra, ..) => EdgeKind::Optional(extra),
1521 Self::Dev(group, ..) => EdgeKind::Dev(group),
1522 }
1523 }
1524}
1525
1526#[derive(Debug, Copy, Clone)]
1527enum RequestedExtras<'env> {
1528 Dependency(&'env BTreeSet<ExtraName>),
1529 Requirement(&'env [ExtraName]),
1530}
1531
1532impl PartialEq for RequestedExtras<'_> {
1533 fn eq(&self, other: &Self) -> bool {
1534 self.iter().eq(other.iter())
1535 }
1536}
1537
1538impl Eq for RequestedExtras<'_> {}
1539
1540impl PartialOrd for RequestedExtras<'_> {
1541 fn partial_cmp(&self, other: &Self) -> Option<Ordering> {
1542 Some(self.cmp(other))
1543 }
1544}
1545
1546impl Ord for RequestedExtras<'_> {
1547 fn cmp(&self, other: &Self) -> Ordering {
1548 self.iter().cmp(other.iter())
1549 }
1550}
1551
1552impl<'env> RequestedExtras<'env> {
1553 fn contains(self, extra: &ExtraName) -> bool {
1554 match self {
1555 Self::Dependency(extras) => extras.contains(extra),
1556 Self::Requirement(extras) => extras.contains(extra),
1557 }
1558 }
1559
1560 fn is_empty(self) -> bool {
1561 match self {
1562 Self::Dependency(extras) => extras.is_empty(),
1563 Self::Requirement(extras) => extras.is_empty(),
1564 }
1565 }
1566
1567 fn iter(self) -> impl Iterator<Item = &'env ExtraName> {
1568 match self {
1569 Self::Dependency(extras) => Either::Left(extras.iter()),
1570 Self::Requirement(extras) => Either::Right(extras.iter()),
1571 }
1572 }
1573}
1574
1575#[derive(Debug, Clone, PartialEq, Eq, Ord, PartialOrd)]
1576enum EdgeKind<'env> {
1577 Prod,
1578 Optional(&'env ExtraName),
1579 Dev(&'env GroupName),
1580}
1581
1582#[derive(Debug, Copy, Clone, PartialEq, Eq, Ord, PartialOrd)]
1584struct Cursor(NodeIndex, Option<EdgeIndex>, UniversalMarker);
1585
1586impl Cursor {
1587 fn new(node: NodeIndex, edge: EdgeIndex, marker: UniversalMarker) -> Self {
1589 Self(node, Some(edge), marker)
1590 }
1591
1592 fn root(node: NodeIndex, marker: UniversalMarker) -> Self {
1594 Self(node, None, marker)
1595 }
1596
1597 fn node(&self) -> NodeIndex {
1599 self.0
1600 }
1601
1602 fn edge(&self) -> Option<EdgeIndex> {
1604 self.1
1605 }
1606
1607 fn marker(&self) -> UniversalMarker {
1609 self.2
1610 }
1611}
1612
1613impl std::fmt::Display for TreeDisplay<'_> {
1614 fn fmt(&self, f: &mut std::fmt::Formatter) -> std::fmt::Result {
1615 use owo_colors::OwoColorize;
1616
1617 let mut deduped = false;
1618 for line in self.render() {
1619 deduped |= line.contains('*');
1620 writeln!(f, "{line}")?;
1621 }
1622
1623 if deduped {
1624 let message = if self.no_dedupe {
1625 "(*) Package tree is a cycle and cannot be shown".italic()
1626 } else {
1627 "(*) Package tree already displayed".italic()
1628 };
1629 writeln!(f, "{message}")?;
1630 }
1631
1632 Ok(())
1633 }
1634}