1use crate::console::{ConsoleOptions, RenderResult, Renderable};
4use crate::segment::Segment;
5use crate::style::Style;
6
7#[derive(Debug, Clone)]
13pub struct TreeGuides {
14 pub space: &'static str,
16 pub continue_line: &'static str,
18 pub fork: &'static str,
20 pub end: &'static str,
22}
23
24pub const ASCII_GUIDES: TreeGuides = TreeGuides {
26 space: " ",
27 continue_line: "| ",
28 fork: "+-- ",
29 end: "`-- ",
30};
31
32pub const TREE_GUIDES: TreeGuides = TreeGuides {
34 space: " ",
35 continue_line: "│ ",
36 fork: "├── ",
37 end: "└── ",
38};
39
40#[derive(Debug, Clone)]
46pub struct Tree {
47 pub label: String,
49 pub style: Style,
51 pub guide_style: Style,
53 pub expanded: bool,
55 pub highlight: bool,
57 pub hide_root: bool,
59 pub children: Vec<Tree>,
61}
62
63impl Tree {
64 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 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 pub fn style(mut self, style: Style) -> Self { self.style = style; self }
86
87 pub fn guide_style(mut self, style: Style) -> Self { self.guide_style = style; self }
89
90 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 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 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}