Add option to draw indentation lines (#8)

Author: nathanwhit

src/lib.rs

@@ -16,12 +16,44 @@ use tracing_subscriber::{
     registry::LookupSpan,
 };
 
+const LINE_VERT: &str = "│";
+const LINE_HORIZ: &str = "─";
+const LINE_BRANCH: &str = "├";
+
+#[derive(Debug)]
+struct Config {
+    ansi: bool,
+    indent_lines: bool,
+    indent_amount: usize,
+}
+
+impl Config {
+    fn with_ansi(self, ansi: bool) -> Self {
+        Self { ansi, ..self }
+    }
+    fn with_indent_lines(self, indent_lines: bool) -> Self {
+        Self {
+            indent_lines,
+            ..self
+        }
+    }
+}
+
+impl Default for Config {
+    fn default() -> Self {
+        Self {
+            ansi: true,
+            indent_lines: false,
+            indent_amount: 2,
+        }
+    }
+}
+
 #[derive(Debug)]
 pub struct HierarchicalLayer {
     stdout: io::Stdout,
-    indent_amount: usize,
-    ansi: bool,
     bufs: Mutex<Buffers>,
+    config: Config,
 }
 
 #[derive(Debug)]
@@ -48,12 +80,13 @@ impl Buffers {
         self.indent_buf.clear();
     }
 
-    fn indent_current(&mut self, indent: usize, indent_amount: usize) {
+    fn indent_current(&mut self, indent: usize, config: &Config) {
         indent_block(
             &mut self.current_buf,
             &mut self.indent_buf,
             indent,
-            indent_amount,
+            config.indent_amount,
+            config.indent_lines,
         );
         self.current_buf.clear();
     }
@@ -102,9 +135,9 @@ impl<'a> Visit for FmtEvent<'a> {
 }
 
 impl<'a> FmtEvent<'a> {
-    fn finish(&mut self, indent: usize, indent_amount: usize) {
+    fn finish(&mut self, indent: usize, config: &Config) {
         self.bufs.current_buf.push('\n');
-        self.bufs.indent_current(indent, indent_amount);
+        self.bufs.indent_current(indent, config);
         self.bufs.flush_indent_buf();
     }
 }
@@ -124,35 +157,111 @@ impl<'a> fmt::Display for ColorLevel<'a> {
     }
 }
 
-fn indent_block(block: &mut String, buf: &mut String, indent: usize, indent_amount: usize) {
-    let lines: Vec<_> = block.lines().collect();
+fn indent_block_with_lines(lines: &[&str], buf: &mut String, indent: usize, indent_amount: usize) {
     let indent_spaces = indent * indent_amount;
-    buf.reserve(block.len() + (lines.len() * indent_spaces));
-    let indent_str = String::from(" ").repeat(indent_spaces);
-    for line in lines {
-        buf.push_str(&indent_str);
+    if lines.len() == 0 {
+        return;
+    } else if indent_spaces == 0 {
+        for line in lines {
+            buf.push_str(line);
+            buf.push('\n');
+        }
+        return;
+    }
+    let mut s = String::with_capacity(indent_spaces);
+
+    // instead of using all spaces to indent, draw a vertical line at every indent level
+    // up until the last indent
+    for i in 0..(indent_spaces - indent_amount) {
+        if i % indent_amount == 0 {
+            s.push_str(LINE_VERT);
+        } else {
+            s.push(' ');
+        }
+    }
+
+    // draw branch
+    buf.push_str(&s);
+    buf.push_str(LINE_BRANCH);
+
+    // add `indent_amount - 1` horizontal lines before the span/event
+    for _ in 0..(indent_amount - 1) {
+        buf.push_str(LINE_HORIZ);
+    }
+    buf.push_str(&lines[0]);
+    buf.push('\n');
+
+    // add the rest of the indentation, since we don't want to draw horizontal lines
+    // for subsequent lines
+    for i in 0..indent_amount {
+        if i % indent_amount == 0 {
+            s.push_str(LINE_VERT);
+        } else {
+            s.push(' ');
+        }
+    }
+
+    // add all of the actual content, with each line preceded by the indent string
+    for line in &lines[1..] {
+        buf.push_str(&s);
         buf.push_str(line);
         buf.push('\n');
     }
 }
 
+fn indent_block(
+    block: &mut String,
+    buf: &mut String,
+    indent: usize,
+    indent_amount: usize,
+    indent_lines: bool,
+) {
+    let lines: Vec<&str> = block.lines().collect();
+    let indent_spaces = indent * indent_amount;
+    buf.reserve(block.len() + (lines.len() * indent_spaces));
+    if indent_lines {
+        indent_block_with_lines(&lines, buf, indent, indent_amount);
+    } else {
+        let indent_str = String::from(" ").repeat(indent_spaces);
+        for line in lines {
+            buf.push_str(&indent_str);
+            buf.push_str(line);
+            buf.push('\n');
+        }
+    }
+}
+
 impl HierarchicalLayer {
     pub fn new(indent_amount: usize) -> Self {
         let ansi = atty::is(atty::Stream::Stdout);
-        Self {
+        let config = Config {
+            ansi,
             indent_amount,
+            ..Default::default()
+        };
+        Self {
             stdout: io::stdout(),
-            ansi,
             bufs: Mutex::new(Buffers::new()),
+            config,
         }
     }
 
     pub fn with_ansi(self, ansi: bool) -> Self {
-        Self { ansi, ..self }
+        Self {
+            config: self.config.with_ansi(ansi),
+            ..self
+        }
+    }
+
+    pub fn with_indent_lines(self, indent_lines: bool) -> Self {
+        Self {
+            config: self.config.with_indent_lines(indent_lines),
+            ..self
+        }
     }
 
     fn styled(&self, style: Style, text: impl AsRef<str>) -> String {
-        if self.ansi {
+        if self.config.ansi {
             style.paint(text.as_ref()).to_string()
         } else {
             text.as_ref().to_string()
@@ -200,7 +309,7 @@ where
         let bufs = &mut *guard;
         let mut current_buf = &mut bufs.current_buf;
 
-        let indent = ctx.scope().count() - 1;
+        let indent = ctx.scope().count().saturating_sub(1);
 
         write!(
             current_buf,
@@ -223,7 +332,7 @@ where
         )
         .unwrap();
 
-        bufs.indent_current(indent, self.indent_amount);
+        bufs.indent_current(indent, &self.config);
         bufs.flush_indent_buf();
         bufs.flush_current_buf(self.stdout.lock());
     }
@@ -271,7 +380,7 @@ where
             .expect("Unable to write to buffer");
         }
         let level = event.metadata().level();
-        let level = if self.ansi {
+        let level = if self.config.ansi {
             ColorLevel(level).to_string()
         } else {
             level.to_string()
@@ -282,9 +391,9 @@ where
             bufs: &mut bufs,
         };
         event.record(&mut visitor);
-        visitor.finish(indent, self.indent_amount);
+        visitor.finish(indent, &self.config);
         bufs.flush_current_buf(self.stdout.lock());
     }
 
-    fn on_close(&self, _id: Id, _ctx: Context<S>) {}
+    fn on_exit(&self, _id: &Id, _ctx: Context<S>) {}
 }