termtree/
lib.rs

1#![cfg_attr(docsrs, feature(doc_auto_cfg))]
2#![warn(clippy::print_stderr)]
3#![warn(clippy::print_stdout)]
4
5#[cfg(test)]
6mod tests;
7
8use std::collections::VecDeque;
9use std::fmt::{self, Display};
10use std::rc::Rc;
11
12/// a simple recursive type which is able to render its
13/// components in a tree-like format
14#[derive(Debug, Clone)]
15pub struct Tree<D: Display> {
16    pub root: D,
17    pub leaves: Vec<Tree<D>>,
18    multiline: bool,
19    glyphs: Option<GlyphPalette>,
20}
21
22impl<D: Display> Tree<D> {
23    pub fn new(root: D) -> Self {
24        Tree {
25            root,
26            leaves: Vec::new(),
27            multiline: false,
28            glyphs: None,
29        }
30    }
31
32    pub fn with_leaves(mut self, leaves: impl IntoIterator<Item = impl Into<Tree<D>>>) -> Self {
33        self.leaves = leaves.into_iter().map(Into::into).collect();
34        self
35    }
36
37    /// Ensure all lines for `root` are indented
38    pub fn with_multiline(mut self, yes: bool) -> Self {
39        self.multiline = yes;
40        self
41    }
42
43    /// Customize the rendering of this node
44    pub fn with_glyphs(mut self, glyphs: GlyphPalette) -> Self {
45        self.glyphs = Some(glyphs);
46        self
47    }
48}
49
50impl<D: Display> Tree<D> {
51    /// Ensure all lines for `root` are indented
52    pub fn set_multiline(&mut self, yes: bool) -> &mut Self {
53        self.multiline = yes;
54        self
55    }
56
57    /// Customize the rendering of this node
58    pub fn set_glyphs(&mut self, glyphs: GlyphPalette) -> &mut Self {
59        self.glyphs = Some(glyphs);
60        self
61    }
62}
63
64impl<D: Display> Tree<D> {
65    pub fn push(&mut self, leaf: impl Into<Tree<D>>) -> &mut Self {
66        self.leaves.push(leaf.into());
67        self
68    }
69}
70
71impl<D: Display> From<D> for Tree<D> {
72    fn from(inner: D) -> Self {
73        Self::new(inner)
74    }
75}
76
77impl<D: Display> Extend<D> for Tree<D> {
78    fn extend<T: IntoIterator<Item = D>>(&mut self, iter: T) {
79        self.leaves.extend(iter.into_iter().map(Into::into));
80    }
81}
82
83impl<D: Display> Extend<Tree<D>> for Tree<D> {
84    fn extend<T: IntoIterator<Item = Tree<D>>>(&mut self, iter: T) {
85        self.leaves.extend(iter);
86    }
87}
88
89impl<D: Display> Display for Tree<D> {
90    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
91        self.root.fmt(f)?; // Pass along `f.alternate()`
92        writeln!(f)?;
93        let mut queue = DisplauQueue::new();
94        let no_space = Rc::new(Vec::new());
95        let default_glyphs = GlyphPalette::new();
96        let glyphs = self.glyphs.as_ref().unwrap_or(&default_glyphs);
97        enqueue_leaves(&mut queue, self, glyphs, no_space);
98        while let Some((last, leaf, glyphs, spaces)) = queue.pop_front() {
99            let mut prefix = (
100                if last {
101                    glyphs.last_item
102                } else {
103                    glyphs.middle_item
104                },
105                glyphs.item_indent,
106            );
107
108            if leaf.multiline {
109                let rest_prefix = (
110                    if last {
111                        glyphs.last_skip
112                    } else {
113                        glyphs.middle_skip
114                    },
115                    glyphs.skip_indent,
116                );
117                debug_assert_eq!(prefix.0.chars().count(), rest_prefix.0.chars().count());
118                debug_assert_eq!(prefix.1.chars().count(), rest_prefix.1.chars().count());
119
120                let root = if f.alternate() {
121                    format!("{:#}", leaf.root)
122                } else {
123                    format!("{:}", leaf.root)
124                };
125                for line in root.lines() {
126                    // print single line
127                    for s in spaces.as_slice() {
128                        s.skip.fmt(f)?;
129                        s.indent.fmt(f)?;
130                    }
131                    prefix.0.fmt(f)?;
132                    prefix.1.fmt(f)?;
133                    line.fmt(f)?;
134                    writeln!(f)?;
135                    prefix = rest_prefix;
136                }
137            } else {
138                // print single line
139                for s in spaces.as_slice() {
140                    s.skip.fmt(f)?;
141                    s.indent.fmt(f)?;
142                }
143                prefix.0.fmt(f)?;
144                prefix.1.fmt(f)?;
145                leaf.root.fmt(f)?; // Pass along `f.alternate()`
146                writeln!(f)?;
147            }
148
149            // recurse
150            if !leaf.leaves.is_empty() {
151                let s: &Vec<SpacePalette> = &spaces;
152                let mut child_spaces = s.clone();
153                child_spaces.push(if last {
154                    glyphs.last_space()
155                } else {
156                    glyphs.middle_space()
157                });
158                let child_spaces = Rc::new(child_spaces);
159                enqueue_leaves(&mut queue, leaf, glyphs, child_spaces);
160            }
161        }
162        Ok(())
163    }
164}
165
166type DisplauQueue<'t, D> = VecDeque<(bool, &'t Tree<D>, &'t GlyphPalette, Rc<Vec<SpacePalette>>)>;
167
168fn enqueue_leaves<'t, D: Display>(
169    queue: &mut DisplauQueue<'t, D>,
170    parent: &'t Tree<D>,
171    parent_glyphs: &'t GlyphPalette,
172    spaces: Rc<Vec<SpacePalette>>,
173) {
174    for (i, leaf) in parent.leaves.iter().rev().enumerate() {
175        let last = i == 0;
176        let glyphs = leaf.glyphs.as_ref().unwrap_or(parent_glyphs);
177        queue.push_front((last, leaf, glyphs, spaces.clone()));
178    }
179}
180
181#[derive(Copy, Clone, Debug, PartialEq, Eq)]
182struct SpacePalette {
183    skip: &'static str,
184    indent: &'static str,
185}
186
187#[derive(Copy, Clone, Debug, PartialEq, Eq)]
188pub struct GlyphPalette {
189    pub middle_item: &'static str,
190    pub last_item: &'static str,
191    pub item_indent: &'static str,
192
193    pub middle_skip: &'static str,
194    pub last_skip: &'static str,
195    pub skip_indent: &'static str,
196}
197
198impl GlyphPalette {
199    pub const fn new() -> Self {
200        Self {
201            middle_item: "├",
202            last_item: "└",
203            item_indent: "── ",
204
205            middle_skip: "│",
206            last_skip: " ",
207            skip_indent: "   ",
208        }
209    }
210
211    fn middle_space(&self) -> SpacePalette {
212        SpacePalette {
213            skip: self.middle_skip,
214            indent: self.skip_indent,
215        }
216    }
217
218    fn last_space(&self) -> SpacePalette {
219        SpacePalette {
220            skip: self.last_skip,
221            indent: self.skip_indent,
222        }
223    }
224}
225
226impl Default for GlyphPalette {
227    fn default() -> Self {
228        Self::new()
229    }
230}