rich/
tree.rs

1//! Trees.
2//!
3//! Port of upstream `rich/tree.py`. A [`Tree`] renders a hierarchy with the
4//! familiar `├──`/`└──` guide lines.
5//!
6//! Slice scope: the default (thin) guides and plain labels. Custom
7//! `guide_style`/label styles and the ASCII/heavy guide sets are deferred with
8//! the rest of `tree.py`.
9
10use crate::cells::cell_len;
11use crate::console::{Console, ConsoleOptions};
12use crate::protocol::Renderable;
13use crate::segment::Segment;
14use crate::style::Style;
15use crate::text::Text;
16use crate::theme::Theme;
17
18// Default (thin) guide segments, matching `TREE_GUIDES[0]`.
19const SPACE: &str = "    ";
20const CONTINUE: &str = "│   ";
21const FORK: &str = "├── ";
22const END: &str = "└── ";
23
24/// A node in a hierarchy. Mirrors `rich.tree.Tree`.
25pub struct Tree {
26    label: String,
27    children: Vec<Tree>,
28}
29
30impl Tree {
31    /// A new tree/subtree with the given label.
32    pub fn new(label: impl Into<String>) -> Self {
33        Tree {
34            label: label.into(),
35            children: Vec::new(),
36        }
37    }
38
39    /// Add a child with `label`, returning a mutable reference to it so further
40    /// descendants can be attached. Mirrors `Tree.add`.
41    pub fn add(&mut self, label: impl Into<String>) -> &mut Tree {
42        self.children.push(Tree::new(label));
43        self.children.last_mut().expect("just pushed a child")
44    }
45
46    /// Recursively render into `lines`. `prefix_first` precedes the label's first
47    /// line; `prefix_rest` precedes wrapped continuation lines and is the base
48    /// for this node's children.
49    fn render_into(
50        &self,
51        theme: &Theme,
52        lines: &mut Vec<Vec<Segment>>,
53        prefix_first: &str,
54        prefix_rest: &str,
55        width: usize,
56    ) {
57        let guide_style = Some(Style::new());
58        let available = width.saturating_sub(cell_len(prefix_first));
59        let mut label_lines =
60            Text::new(&self.label).render_lines(theme, &Style::new(), Some(available));
61        if label_lines.is_empty() {
62            label_lines.push(Vec::new());
63        }
64
65        for (index, label_line) in label_lines.into_iter().enumerate() {
66            let prefix = if index == 0 {
67                prefix_first
68            } else {
69                prefix_rest
70            };
71            let mut line = Vec::new();
72            if !prefix.is_empty() {
73                line.push(Segment::new(prefix.to_string(), guide_style.clone()));
74            }
75            line.extend(label_line);
76            lines.push(line);
77        }
78
79        let last_index = self.children.len().saturating_sub(1);
80        for (index, child) in self.children.iter().enumerate() {
81            let last = index == last_index;
82            let child_first = format!("{prefix_rest}{}", if last { END } else { FORK });
83            let child_rest = format!("{prefix_rest}{}", if last { SPACE } else { CONTINUE });
84            child.render_into(theme, lines, &child_first, &child_rest, width);
85        }
86    }
87}
88
89impl Renderable for Tree {
90    fn rich_render(&self, console: &Console, options: &ConsoleOptions) -> Vec<Segment> {
91        let mut lines: Vec<Vec<Segment>> = Vec::new();
92        self.render_into(console.theme(), &mut lines, "", "", options.max_width);
93
94        let mut segments = Vec::new();
95        let last = lines.len().saturating_sub(1);
96        for (index, line) in lines.into_iter().enumerate() {
97            segments.extend(line);
98            if index != last {
99                segments.push(Segment::line());
100            }
101        }
102        segments
103    }
104}
105
106#[cfg(test)]
107mod tests {
108    use super::*;
109    use crate::color::ColorSystem;
110
111    fn console() -> Console {
112        Console::builder()
113            .force_terminal(true)
114            .color_system(Some(ColorSystem::Truecolor))
115            .width(40)
116            .build()
117    }
118
119    #[test]
120    fn nested_tree() {
121        let mut tree = Tree::new("root");
122        let a = tree.add("child A");
123        a.add("leaf A1");
124        a.add("leaf A2");
125        tree.add("child B");
126        let out = console().render_export(&tree);
127        let expected = concat!(
128            "root\n",
129            "├── child A\n",
130            "│   ├── leaf A1\n",
131            "│   └── leaf A2\n",
132            "└── child B\n",
133        );
134        assert_eq!(out, expected);
135    }
136}
rich/tree.rs

rich/
tree.rs
