Skip to main content

rusty_rich/
tree.rs

1//! Tree — hierarchical tree rendering. Equivalent to Rich's `tree.py`.
2
3use crate::console::{ConsoleOptions, RenderResult, Renderable};
4use crate::segment::Segment;
5use crate::style::Style;
6
7// ---------------------------------------------------------------------------
8// Guide types
9// ---------------------------------------------------------------------------
10
11/// The characters used for tree guide lines at each position.
12#[derive(Debug, Clone)]
13pub struct TreeGuides {
14    /// Space guide (no line).
15    pub space: &'static str,
16    /// Continue guide (vertical line).
17    pub continue_line: &'static str,
18    /// Fork guide (branch with more siblings).
19    pub fork: &'static str,
20    /// End guide (last sibling).
21    pub end: &'static str,
22}
23
24/// ASCII-only guides.
25pub const ASCII_GUIDES: TreeGuides = TreeGuides {
26    space: "    ",
27    continue_line: "|   ",
28    fork: "+-- ",
29    end: "`-- ",
30};
31
32/// Default Unicode guides (like Rich's `TREE_GUIDES[0]`).
33pub const TREE_GUIDES: TreeGuides = TreeGuides {
34    space: "    ",
35    continue_line: "│   ",
36    fork: "├── ",
37    end: "└── ",
38};
39
40// ---------------------------------------------------------------------------
41// Tree
42// ---------------------------------------------------------------------------
43
44/// A renderable for a tree structure.
45#[derive(Debug, Clone)]
46pub struct Tree {
47    /// The label for this node.
48    pub label: String,
49    /// Style for this node's label.
50    pub style: Style,
51    /// Style for the guide lines.
52    pub guide_style: Style,
53    /// If true, children are visible.
54    pub expanded: bool,
55    /// If true, highlight string labels.
56    pub highlight: bool,
57    /// If true, don't show the root node.
58    pub hide_root: bool,
59    /// Children of this node.
60    pub children: Vec<Tree>,
61}
62
63impl Tree {
64    /// Create a new tree node with the given label.
65    pub fn new(label: impl Into<String>) -> Self {
66        Self {
67            label: label.into(),
68            style: Style::new(),
69            guide_style: Style::new(),
70            expanded: true,
71            highlight: false,
72            hide_root: false,
73            children: Vec::new(),
74        }
75    }
76
77    /// Add a child node, returning a mutable reference to it.
78    pub fn add(&mut self, label: impl Into<String>) -> &mut Tree {
79        let child = Tree::new(label);
80        self.children.push(child);
81        self.children.last_mut().unwrap()
82    }
83
84    /// Set the style for this node.
85    pub fn style(mut self, style: Style) -> Self { self.style = style; self }
86
87    /// Set the guide style.
88    pub fn guide_style(mut self, style: Style) -> Self { self.guide_style = style; self }
89
90    /// Hide the root node.
91    pub fn hide_root(mut self) -> Self { self.hide_root = true; self }
92}
93
94impl Renderable for Tree {
95    fn render(&self, options: &ConsoleOptions) -> RenderResult {
96        let guides = if options.ascii_only { &ASCII_GUIDES } else { &TREE_GUIDES };
97        let mut lines: Vec<Vec<Segment>> = Vec::new();
98
99        if !self.hide_root {
100            lines.push(vec![Segment::new(&self.label), Segment::line()]);
101        }
102
103        let last_idx = self.children.len().saturating_sub(1);
104        for (i, child) in self.children.iter().enumerate() {
105            let is_last = i == last_idx;
106            self.render_node(child, &mut lines, guides, "", is_last, options);
107        }
108
109        RenderResult { lines, items: Vec::new() }
110    }
111}
112
113impl Tree {
114    fn render_node(
115        &self,
116        node: &Tree,
117        lines: &mut Vec<Vec<Segment>>,
118        guides: &TreeGuides,
119        prefix: &str,
120        is_last: bool,
121        options: &ConsoleOptions,
122    ) {
123        let connector = if is_last { guides.end } else { guides.fork };
124        let guide_ansi = self.guide_style.to_ansi();
125        let guide_reset = if guide_ansi.is_empty() { "" } else { "\x1b[0m" };
126
127        // Render this node
128        let guide_str = format!("{prefix}{connector}");
129        lines.push(vec![
130            Segment::new(format!("{guide_ansi}{guide_str}{guide_reset}")),
131            Segment::new(&node.label),
132            Segment::line(),
133        ]);
134
135        // Children continuation prefix
136        let child_prefix = if is_last {
137            format!("{prefix}{}", guides.space)
138        } else {
139            format!("{prefix}{}", guides.continue_line)
140        };
141        let child_prefix_styled = format!("{guide_ansi}{child_prefix}{guide_reset}");
142
143        let last_child = node.children.len().saturating_sub(1);
144        for (i, child) in node.children.iter().enumerate() {
145            let child_is_last = i == last_child;
146            self.render_node(child, lines, guides, &child_prefix_styled, child_is_last, options);
147        }
148    }
149}
150
151#[cfg(test)]
152mod tests {
153    use super::*;
154    use crate::console::ConsoleOptions;
155
156    #[test]
157    fn test_simple_tree() {
158        let mut tree = Tree::new("Root");
159        tree.add("Child 1");
160        tree.add("Child 2");
161
162        let opts = ConsoleOptions::default();
163        let result = tree.render(&opts);
164        let ansi = result.to_ansi();
165        assert!(ansi.contains("Root"));
166        assert!(ansi.contains("Child 1"));
167        assert!(ansi.contains("Child 2"));
168    }
169
170    #[test]
171    fn test_nested_tree() {
172        let mut tree = Tree::new("Root");
173        let child = tree.add("A");
174        child.add("A.1");
175        tree.add("B");
176
177        let opts = ConsoleOptions::default();
178        let result = tree.render(&opts);
179        let ansi = result.to_ansi();
180        assert!(ansi.contains("A.1"));
181    }
182}