prettier print with indentation object

Author: LoZack19

src/display.rs

@@ -0,0 +1,86 @@
+use std::fmt::Display;
+
+/// Indentation token
+#[derive(Debug)]
+struct Token {
+    /// Is followed by a brother
+    siblings: bool,
+    /// Is intermediate while printing children
+    children: bool,
+}
+
+impl ToString for Token {
+    fn to_string(&self) -> String {
+        let Token { siblings, children } = self;
+
+        match (siblings, children) {
+            (true, true) => "│   ",
+            (true, false) => "├── ",
+            (false, true) => "    ",
+            (false, false) => "└── ",
+        }
+        .to_string()
+    }
+}
+
+impl Token {
+    /// Create a new indentation token
+    fn new(siblings: bool) -> Self {
+        Token {
+            siblings,
+            children: false,
+        }
+    }
+
+    /// Set children flag before starting displaying children
+    fn set_children(&mut self) {
+        self.children = true;
+    }
+}
+
+/// Manages the state during the display operation
+#[derive(Debug)]
+pub struct Indentation {
+    tokens: Vec<Token>,
+    ignore_root: bool,
+}
+
+impl Display for Indentation {
+    fn fmt(&self, f: &mut std::fmt::Formatter) -> std::fmt::Result {
+        let first: usize = if self.ignore_root { 1 } else { 0 };
+
+        for token in &self.tokens[first..] {
+            write!(f, "{}", token.to_string())?;
+        }
+
+        Ok(())
+    }
+}
+
+impl Indentation {
+    /// Creates a new indentation handler
+    pub fn new(ignore_root: bool) -> Self {
+        Indentation {
+            tokens: Vec::new(),
+            ignore_root,
+        }
+    }
+
+    /// Adds a new layer of indentation
+    pub fn indent(&mut self, siblings: bool) -> &mut Self {
+        // Setup children mode for previous tokens
+        let len = self.tokens.len();
+        if len > 0 {
+            self.tokens[len - 1].set_children();
+        }
+
+        self.tokens.push(Token::new(siblings));
+        self
+    }
+
+    /// Removes the last layer of indentation
+    pub fn deindent(&mut self) -> &mut Self {
+        self.tokens.pop();
+        self
+    }
+}

src/lib.rs

@@ -719,36 +719,34 @@ impl<T: Debug> Debug for Tree<T> {
     }
 }
 
+// Handles display
+mod display;
+
 impl<T: Display> Display for Tree<T> {
     fn fmt(&self, f: &mut Formatter) -> Result<(), fmt::Error> {
+        use display::Indentation;
         use iter::Edge;
 
-        fn indent(f: &mut Formatter, level: usize) -> Result<(), fmt::Error> {
-            for _ in 0..level {
-                write!(f, "    ")?;
-            }
-            Ok(())
-        }
-
-        let mut level: usize = 0;
+        let mut indent: Indentation = Indentation::new(true);
 
         for edge in self.root().traverse() {
             match edge {
                 Edge::Open(node) if node.has_children() => {
-                    indent(f, level)?;
-                    write!(f, "{}\n", node.value())?;
-                    level += 1;
+                    indent.indent(node.next_sibling().is_some());
+                    write!(f, "{}{}\n", indent.to_string(), node.value())?;
                 }
                 Edge::Open(node) if node.next_sibling().is_some() => {
-                    indent(f, level)?;
-                    write!(f, "{}\n", node.value())?;
+                    indent.indent(node.next_sibling().is_some());
+                    write!(f, "{}{}\n", indent.to_string(), node.value())?;
+                    indent.deindent();
                 }
                 Edge::Open(node) => {
-                    indent(f, level)?;
-                    write!(f, "{}\n", node.value())?;
+                    indent.indent(node.next_sibling().is_some());
+                    write!(f, "{}{}\n", indent.to_string(), node.value())?;
+                    indent.deindent();
                 }
                 Edge::Close(node) if node.has_children() => {
-                    level -= 1;
+                    indent.deindent();
                 }
                 _ => {}
             }

tests/tree.rs

@@ -171,7 +171,7 @@ fn test_display() {
     };
 
     let repr = format!("{tree}");
-    let expected = "root\n    a\n        child 1\n    b\n        child 2\n";
+    let expected = "root\n├── a\n│   └── child 1\n└── b\n    └── child 2\n";
 
     assert_eq!(repr, expected);
 
@@ -184,7 +184,7 @@ fn test_display() {
     };
 
     let repr = format!("{tree}");
-    let expected = "root\n    a\n    b\n        x\n        y\n    c\n";
+    let expected = "root\n├── a\n├── b\n│   ├── x\n│   └── y\n└── c\n";
 
     assert_eq!(repr, expected);
 }