1mod format;
2mod macros;
3
4use std::error::Error;
5use std::fmt::{self, Arguments, Display, Write};
6
7pub use format::*;
8
9type Result<T> = std::result::Result<T, Box<dyn Error>>;
10
11pub trait TreeDisplay {
12 fn fmt(&self, tf: &mut TreeFormatter<'_>) -> fmt::Result;
13}
14
15pub struct TreeFormatter<'a> {
16 inner: &'a mut (dyn Write + 'a),
17
18 level: usize,
19 current_context: String,
20 context_indices: Vec<usize>,
21 context_format: ContextFormat<'a>,
22 prefix_format: ItemPrefixFormat<'a>,
23}
24
25impl Write for TreeFormatter<'_> {
26 #[inline]
27 fn write_str(&mut self, s: &str) -> fmt::Result {
28 self.inner.write_str(s)
29 }
30}
31
32impl<'a> TreeFormatter<'a> {
33 const DEFAULT_BRANCH_CAPACITY: usize = 3;
34
35 pub fn new(title: impl Display, inner: &'a mut impl Write) -> Result<Self> {
36 Self::with_context(title, inner, "")
37 }
38
39 pub fn with_context(title: impl Display, inner: &'a mut impl Write, context: &'a str) -> Result<Self> {
40 writeln!(inner, "{context}{title}")?;
41 Ok(Self {
42 inner,
43 level: 0,
44 current_context: String::from(context),
45 context_indices: Vec::with_capacity(Self::DEFAULT_BRANCH_CAPACITY),
46 context_format: ContextFormat::default(),
47 prefix_format: ItemPrefixFormat::default(),
48 })
49 }
50
51 pub fn context_format(&mut self, context_format: ContextFormat<'a>) -> &mut Self {
52 self.context_format = context_format;
53 self
54 }
55
56 pub fn prefix_format(&mut self, prefix_format: ItemPrefixFormat<'a>) -> &mut Self {
57 self.prefix_format = prefix_format;
58 self
59 }
60
61 pub fn begin_level(&mut self, is_last: bool) {
62 self.level += 1;
63 if self.level > 1 {
64 self.context_indices.push(self.current_context.len());
65 self.current_context.push_str(self.context_format.context(is_last));
66 }
67 }
68
69 pub fn end_level(&mut self) {
70 self.level -= 1;
71 if let Some(last_idx) = self.context_indices.pop() {
72 self.current_context.truncate(last_idx);
73 }
74 }
75
76 pub fn write(&mut self, is_last: bool, item: impl Display) -> fmt::Result {
77 writeln!(self.inner, "{}{}{}", self.current_context, self.prefix_format.prefix(is_last), item)?;
78 if is_last && self.context_indices.len() > self.level - 1 {
79 let last_idx = self.context_indices.len() - 1;
80 self
81 .current_context
82 .replace_range(self.context_indices[last_idx].., self.context_format.empty);
83 }
84 Ok(())
85 }
86
87 pub fn write_fmt(&mut self, is_last: bool, item: Arguments<'a>) -> fmt::Result {
88 self.write(is_last, item)
89 }
90
91 pub fn write_level(&mut self, is_last: bool, items: impl Iterator<Item = impl Display>) -> fmt::Result {
92 self.begin_level(is_last);
93 let mut items = items.peekable();
94 while let Some(item) = items.next() {
95 self.write(items.peek().is_none(), item)?;
96 }
97 self.end_level();
98 Ok(())
99 }
100
101 pub fn write_tree(&mut self, tree: impl TreeDisplay) -> fmt::Result {
102 tree.fmt(self)
103 }
104}
105
106#[cfg(test)]
107mod tests {
108 use super::*;
109 use indoc::indoc;
110
111 fn generate_tree(fmt: &mut TreeFormatter) -> Result<()> {
112 tree_indent!(fmt);
113 tree_write!(fmt, "Child Level 1 #1")?;
114 fmt.write_level(false, (1..=3u8).map(|x| format!("Child Level 2 #{x}")))?;
115 tree_write!(fmt, "Child Level 1 #2")?;
116 tree_indent!(fmt);
117 tree_write!(fmt, "Child Level 2 #4")?;
118 tree_write_last!(fmt, "Child Level 2 #5")?;
119 tree_indent_last!(fmt);
120 tree_write!(fmt, "Child Level 3 #1")?;
121 tree_indent!(fmt);
122 tree_write!(fmt, "Child Level 4 #1")?;
123 tree_write_last!(fmt, "Child Level 4 #2")?;
124 tree_unindent!(fmt);
125 tree_write_last!(fmt, "Child Level 3 #2")?;
126 tree_unindent!(fmt);
127 tree_unindent!(fmt);
128 tree_write_last!(fmt, "Child Level 1 #3")?;
129 tree_indent_last!(fmt);
130 tree_write!(fmt, "Child Level 2 #6")?;
131 tree_write_last!(fmt, "Child Level 2 #7")?;
132 fmt.write_level(true, (3..=5u8).map(|x| format!("Child Level 3 #{x}")))?;
133 tree_unindent!(fmt);
134 Ok(())
135 }
136
137 #[test]
138 fn test_tree_formatter() {
139 let mut buf = String::new();
140 let mut formatter = TreeFormatter::new("Root", &mut buf).unwrap();
141
142 generate_tree(&mut formatter).unwrap();
143 assert_eq!(
144 buf,
145 indoc! {"
146 Root
147 ├── Child Level 1 #1
148 │ ├── Child Level 2 #1
149 │ ├── Child Level 2 #2
150 │ └── Child Level 2 #3
151 ├── Child Level 1 #2
152 │ ├── Child Level 2 #4
153 │ └── Child Level 2 #5
154 │ ├── Child Level 3 #1
155 │ │ ├── Child Level 4 #1
156 │ │ └── Child Level 4 #2
157 │ └── Child Level 3 #2
158 └── Child Level 1 #3
159 ├── Child Level 2 #6
160 └── Child Level 2 #7
161 ├── Child Level 3 #3
162 ├── Child Level 3 #4
163 └── Child Level 3 #5
164 "}
165 );
166 }
167
168 #[test]
169 fn test_tree_formatter_custom_context_format() {
170 let mut buf = String::new();
171 let mut formatter = TreeFormatter::new("Root", &mut buf).unwrap();
172
173 formatter.context_format(ContextFormat::new("* * ", "║ | "));
174
175 generate_tree(&mut formatter).unwrap();
176 assert_eq!(
177 buf,
178 indoc! {"
179 Root
180 ├── Child Level 1 #1
181 ║ | ├── Child Level 2 #1
182 ║ | ├── Child Level 2 #2
183 ║ | └── Child Level 2 #3
184 ├── Child Level 1 #2
185 ║ | ├── Child Level 2 #4
186 ║ | └── Child Level 2 #5
187 ║ | * * ├── Child Level 3 #1
188 ║ | * * ║ | ├── Child Level 4 #1
189 ║ | * * ║ | └── Child Level 4 #2
190 ║ | * * └── Child Level 3 #2
191 └── Child Level 1 #3
192 * * ├── Child Level 2 #6
193 * * └── Child Level 2 #7
194 * * * * ├── Child Level 3 #3
195 * * * * ├── Child Level 3 #4
196 * * * * └── Child Level 3 #5
197 "}
198 )
199 }
200
201 #[test]
202 fn test_tree_formatter_custom_prefix_format() {
203 let mut buf = String::new();
204 let mut formatter = TreeFormatter::new("Root", &mut buf).unwrap();
205 formatter.context_format(ContextFormat::new("* * ", "║ | "));
206 formatter.prefix_format(ItemPrefixFormat::new("╠═══", "╚═══"));
207
208 generate_tree(&mut formatter).unwrap();
209 assert_eq!(
210 buf,
211 indoc! {"
212 Root
213 ╠═══Child Level 1 #1
214 ║ | ╠═══Child Level 2 #1
215 ║ | ╠═══Child Level 2 #2
216 ║ | ╚═══Child Level 2 #3
217 ╠═══Child Level 1 #2
218 ║ | ╠═══Child Level 2 #4
219 ║ | ╚═══Child Level 2 #5
220 ║ | * * ╠═══Child Level 3 #1
221 ║ | * * ║ | ╠═══Child Level 4 #1
222 ║ | * * ║ | ╚═══Child Level 4 #2
223 ║ | * * ╚═══Child Level 3 #2
224 ╚═══Child Level 1 #3
225 * * ╠═══Child Level 2 #6
226 * * ╚═══Child Level 2 #7
227 * * * * ╠═══Child Level 3 #3
228 * * * * ╠═══Child Level 3 #4
229 * * * * ╚═══Child Level 3 #5
230 "}
231 )
232 }
233
234 #[test]
235 fn test_tree_formatter_with_context() {
236 let mut buf = String::new();
237 let mut formatter = TreeFormatter::with_context("Root", &mut buf, "Some context: ").unwrap();
238
239 generate_tree(&mut formatter).unwrap();
240 assert_eq!(
241 buf,
242 indoc! {"
243 Some context: Root
244 Some context: ├── Child Level 1 #1
245 Some context: │ ├── Child Level 2 #1
246 Some context: │ ├── Child Level 2 #2
247 Some context: │ └── Child Level 2 #3
248 Some context: ├── Child Level 1 #2
249 Some context: │ ├── Child Level 2 #4
250 Some context: │ └── Child Level 2 #5
251 Some context: │ ├── Child Level 3 #1
252 Some context: │ │ ├── Child Level 4 #1
253 Some context: │ │ └── Child Level 4 #2
254 Some context: │ └── Child Level 3 #2
255 Some context: └── Child Level 1 #3
256 Some context: ├── Child Level 2 #6
257 Some context: └── Child Level 2 #7
258 Some context: ├── Child Level 3 #3
259 Some context: ├── Child Level 3 #4
260 Some context: └── Child Level 3 #5
261 "}
262 )
263 }
264}