dmntk_common/
ascii_tree.rs

1//! # ASCII tree with colors
2
3use crate::{ColorMode, ASCII_RESET};
4use std::fmt;
5use std::fmt::{Display, Formatter, Write};
6
7/// Text with associated color control sequence to be displayed in terminal.
8pub struct AsciiText {
9  color: Option<String>,
10  text: String,
11}
12
13impl AsciiText {
14  /// Creates a new text without color control sequence.
15  pub fn new(text: &str) -> Self {
16    Self {
17      color: None,
18      text: text.to_string(),
19    }
20  }
21
22  /// Creates a new text with associated color control sequence.
23  pub fn with_color(text: &str, color: &str) -> Self {
24    Self {
25      color: Some(color.to_string()),
26      text: text.to_string(),
27    }
28  }
29
30  /// Returns a text with color control sequence based on provided [ColorMode].
31  pub fn mode(&self, color_mode: &ColorMode) -> Self {
32    AsciiText {
33      color: match color_mode {
34        ColorMode::On => self.color.clone(),
35        _ => None,
36      },
37      text: self.text.clone(),
38    }
39  }
40}
41
42impl Display for AsciiText {
43  /// Prints a text with associated color control sequence.
44  fn fmt(&self, f: &mut Formatter<'_>) -> fmt::Result {
45    if let Some(color) = &self.color {
46      write!(f, "{}{}{}", color, self.text, ASCII_RESET)
47    } else {
48      write!(f, "{}", self.text)
49    }
50  }
51}
52
53/// Collection of [AsciiText] segments with associated color control sequence.
54pub struct AsciiLine(Vec<AsciiText>);
55
56impl AsciiLine {
57  pub fn builder() -> AsciiLineBuilder {
58    AsciiLineBuilder(vec![])
59  }
60
61  pub fn mode(&self, color_mode: &ColorMode) -> Self {
62    AsciiLine(self.0.iter().map(|ascii_text| ascii_text.mode(color_mode)).collect())
63  }
64}
65
66impl Display for AsciiLine {
67  fn fmt(&self, f: &mut Formatter<'_>) -> fmt::Result {
68    for text in &self.0 {
69      write!(f, "{}", text)?
70    }
71    Ok(())
72  }
73}
74
75/// Builder for [AsciiLine].
76pub struct AsciiLineBuilder(Vec<AsciiText>);
77
78impl AsciiLineBuilder {
79  pub fn text(mut self, text: &str) -> Self {
80    self.0.push(AsciiText::new(text));
81    self
82  }
83
84  pub fn with_color(mut self, text: &str, color: &str) -> Self {
85    self.0.push(AsciiText::with_color(text, color));
86    self
87  }
88
89  pub fn indent(mut self) -> Self {
90    self.0.push(AsciiText::new("  "));
91    self
92  }
93
94  pub fn colon(mut self) -> Self {
95    self.0.push(AsciiText::new(":"));
96    self
97  }
98
99  pub fn colon_space(mut self) -> Self {
100    self.0.push(AsciiText::new(": "));
101    self
102  }
103
104  pub fn build(self) -> AsciiLine {
105    AsciiLine(self.0)
106  }
107}
108
109/// Types of nodes in the coloured ASCII tree.
110pub enum AsciiNode {
111  /// Intermediary (or root) node in the tree (always has child nodes).
112  Node(AsciiLine, Vec<AsciiNode>),
113  /// Node being the leaf in the tree (has no child nodes).
114  Leaf(Vec<AsciiLine>),
115}
116
117impl AsciiNode {
118  pub fn leaf_builder() -> AsciiLeafBuilder {
119    AsciiLeafBuilder(vec![])
120  }
121
122  pub fn node_builder(line: AsciiLine) -> AsciiNodeBuilder {
123    AsciiNodeBuilder(line, vec![])
124  }
125}
126
127/// Builder for [AsciiNode::Leaf].
128pub struct AsciiLeafBuilder(Vec<AsciiLine>);
129
130impl AsciiLeafBuilder {
131  pub fn line(mut self, line: AsciiLine) -> Self {
132    self.0.push(line);
133    self
134  }
135
136  pub fn add_line(&mut self, line: AsciiLine) {
137    self.0.push(line);
138  }
139
140  pub fn build(self) -> AsciiNode {
141    AsciiNode::Leaf(self.0)
142  }
143}
144
145/// Builder for [AsciiNode::Node].
146pub struct AsciiNodeBuilder(AsciiLine, Vec<AsciiNode>);
147
148impl AsciiNodeBuilder {
149  pub fn child(mut self, child: AsciiNode) -> Self {
150    self.1.push(child);
151    self
152  }
153
154  pub fn opt_child(mut self, opt_child: Option<AsciiNode>) -> Self {
155    if let Some(child) = opt_child {
156      self.1.push(child);
157    }
158    self
159  }
160
161  pub fn add_child(&mut self, child: AsciiNode) {
162    self.1.push(child);
163  }
164
165  pub fn build(self) -> AsciiNode {
166    AsciiNode::Node(self.0, self.1)
167  }
168}
169
170/// Writes the tree to provided writer starting from specified node.
171pub fn write(f: &mut dyn Write, node: &AsciiNode, color_mode: &ColorMode) -> fmt::Result {
172  write_node(f, node, &[], color_mode)
173}
174
175/// Writes the tree to provided writer starting from specified node with indentation.
176pub fn write_indented(f: &mut dyn Write, node: &AsciiNode, color_mode: &ColorMode, indent: usize) -> fmt::Result {
177  let mut buffer = String::new();
178  let _ = write_node(&mut buffer, node, &[], color_mode);
179  let indent = " ".repeat(indent);
180  let mut tree = String::new();
181  for line in buffer.lines() {
182    let _ = writeln!(&mut tree, "{}{}", indent, line);
183  }
184  write!(f, "{}", tree)
185}
186
187/// Writes the tree node to provided writer.
188fn write_node(f: &mut dyn Write, tree: &AsciiNode, level: &[usize], color_mode: &ColorMode) -> fmt::Result {
189  const NONE: &str = "   ";
190  const EDGE: &str = " └─";
191  const PIPE: &str = " │ ";
192  const FORK: &str = " ├─";
193
194  let max_pos = level.len();
195  let mut second_line = String::new();
196  for (pos, lev) in level.iter().enumerate() {
197    let last_row = pos == max_pos - 1;
198    if *lev == 1 {
199      if !last_row {
200        write!(f, "{}", NONE)?
201      } else {
202        write!(f, "{}", EDGE)?
203      }
204      second_line.push_str(NONE);
205    } else {
206      if !last_row {
207        write!(f, "{}", PIPE)?
208      } else {
209        write!(f, "{}", FORK)?
210      }
211      second_line.push_str(PIPE);
212    }
213  }
214  match tree {
215    AsciiNode::Node(title, children) => {
216      let mut deep = children.len();
217      writeln!(f, " {}", title.mode(color_mode))?;
218      for node in children {
219        let mut level_next = level.to_vec();
220        level_next.push(deep);
221        deep -= 1;
222        write_node(f, node, &level_next, color_mode)?;
223      }
224    }
225    AsciiNode::Leaf(lines) => {
226      for (i, line) in lines.iter().enumerate() {
227        match i {
228          0 => writeln!(f, " {}", line.mode(color_mode))?,
229          _ => writeln!(f, "{} {}", second_line, line.mode(color_mode))?,
230        }
231      }
232    }
233  }
234  Ok(())
235}
236
237#[cfg(test)]
238mod tests {
239  use super::*;
240
241  const EXPECTED: &str = r#"
242 node 4
243 ├─ node 1
244 │  ├─ line 1_1
245 │  │  line 1_2
246 │  │  line 1_3
247 │  │  line 1_4
248 │  └─ only one line
249 ├─ node 2
250 │  ├─ only one line
251 │  ├─ line 2_1
252 │  │  line 2_2
253 │  │  line 2_3
254 │  │  line 2_4
255 │  └─ only one line
256 └─ node 3
257    ├─ node 1
258    │  ├─ line 3_1_1
259    │  │  line 3_1_2
260    │  │  line 3_1_3
261    │  │  line 3_1_4
262    │  └─ only one line
263    ├─ line 3_1
264    │  line 3_2
265    │  line 3_3
266    │  line 3_4
267    └─ only one line
268"#;
269
270  #[test]
271  fn test_ascii_tree() {
272    let root = AsciiNode::node_builder(AsciiLine::builder().text("node 4").build())
273      .child(
274        AsciiNode::node_builder(AsciiLine::builder().text("node 1").build())
275          .child(
276            AsciiNode::leaf_builder()
277              .line(AsciiLine::builder().text("line 1_1").build())
278              .line(AsciiLine::builder().text("line 1_2").build())
279              .line(AsciiLine::builder().text("line 1_3").build())
280              .line(AsciiLine::builder().text("line 1_4").build())
281              .build(),
282          )
283          .child(AsciiNode::leaf_builder().line(AsciiLine::builder().text("only one line").build()).build())
284          .build(),
285      )
286      .child(
287        AsciiNode::node_builder(AsciiLine::builder().text("node 2").build())
288          .child(AsciiNode::leaf_builder().line(AsciiLine::builder().text("only one line").build()).build())
289          .child(
290            AsciiNode::leaf_builder()
291              .line(AsciiLine::builder().text("line 2_1").build())
292              .line(AsciiLine::builder().text("line 2_2").build())
293              .line(AsciiLine::builder().text("line 2_3").build())
294              .line(AsciiLine::builder().text("line 2_4").build())
295              .build(),
296          )
297          .child(AsciiNode::leaf_builder().line(AsciiLine::builder().text("only one line").build()).build())
298          .build(),
299      )
300      .child(
301        AsciiNode::node_builder(AsciiLine::builder().text("node 3").build())
302          .child(
303            AsciiNode::node_builder(AsciiLine::builder().text("node 1").build())
304              .child(
305                AsciiNode::leaf_builder()
306                  .line(AsciiLine::builder().text("line 3_1_1").build())
307                  .line(AsciiLine::builder().text("line 3_1_2").build())
308                  .line(AsciiLine::builder().text("line 3_1_3").build())
309                  .line(AsciiLine::builder().text("line 3_1_4").build())
310                  .build(),
311              )
312              .child(AsciiNode::leaf_builder().line(AsciiLine::builder().text("only one line").build()).build())
313              .build(),
314          )
315          .child(
316            AsciiNode::leaf_builder()
317              .line(AsciiLine::builder().text("line 3_1").build())
318              .line(AsciiLine::builder().text("line 3_2").build())
319              .line(AsciiLine::builder().text("line 3_3").build())
320              .line(AsciiLine::builder().text("line 3_4").build())
321              .build(),
322          )
323          .child(AsciiNode::leaf_builder().line(AsciiLine::builder().text("only one line").build()).build())
324          .build(),
325      )
326      .build();
327
328    let mut output = String::new();
329    let _ = writeln!(&mut output);
330    let _ = write(&mut output, &root, &ColorMode::Off);
331    assert_eq!(EXPECTED, output);
332  }
333}
dmntk_common/ascii_tree.rs

dmntk_common/
ascii_tree.rs
