From 5081bf0d085fa50553d1ca27ce94cb46fcfea603 Mon Sep 17 00:00:00 2001 From: David Barsky Date: Sat, 11 Jul 2020 17:13:25 -0400 Subject: [PATCH] Add a `MakeWriter` param to `HierarchicalLayer` (#10) Authored-by: David Barsky Co-authored-by: Nathan Whitaker <17734409+nathanwhit@users.noreply.github.com> --- Cargo.toml | 4 +- examples/basic.rs | 6 +- examples/stderr.rs | 40 +++++++ src/format.rs | 196 ++++++++++++++++++++++++++++++++++ src/lib.rs | 258 +++++++++++---------------------------------- 5 files changed, 304 insertions(+), 200 deletions(-) create mode 100644 examples/stderr.rs create mode 100644 src/format.rs diff --git a/Cargo.toml b/Cargo.toml index 1f3d636..32916a2 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -1,6 +1,6 @@ [package] name = "tracing-tree" -version = "0.1.2" +version = "0.1.3" authors = ["David Barsky ", "Nathan Whitaker"] edition = "2018" license = "MIT OR Apache-2.0" @@ -10,7 +10,7 @@ description = "a tracing layer that represents prints a heirarchal tree of spans [dependencies] tracing = "0.1" -tracing-subscriber = "0.2" +tracing-subscriber = { version = "0.2", default-features = false, features = ["registry", "fmt"] } quanta = "0.3.1" termcolor = "1.0.5" ansi_term = "0.12.1" diff --git a/examples/basic.rs b/examples/basic.rs index eee7775..e533f0e 100644 --- a/examples/basic.rs +++ b/examples/basic.rs @@ -3,7 +3,11 @@ use tracing_subscriber::{layer::SubscriberExt, registry::Registry}; use tracing_tree::HierarchicalLayer; fn main() { - let subscriber = Registry::default().with(HierarchicalLayer::new(2).with_indent_lines(true)); + let layer = HierarchicalLayer::default() + .with_indent_lines(true) + .with_indent_amount(2); + + let subscriber = Registry::default().with(layer); tracing::subscriber::set_global_default(subscriber).unwrap(); let app_span = span!(Level::TRACE, "hierarchical-example", version = %0.1); diff --git a/examples/stderr.rs b/examples/stderr.rs new file mode 100644 index 0000000..071048c --- /dev/null +++ b/examples/stderr.rs @@ -0,0 +1,40 @@ +use tracing::{debug, info, instrument}; +use tracing_subscriber::{layer::SubscriberExt, registry::Registry}; +use tracing_tree::HierarchicalLayer; + +#[instrument] +fn nth_fibonacci(n: u64) -> u64 { + if n == 0 || n == 1 { + debug!("Base case"); + 1 + } else { + debug!("Recursing"); + nth_fibonacci(n - 1) + nth_fibonacci(n - 2) + } +} + +#[instrument] +fn fibonacci_seq(to: u64) -> Vec { + let mut sequence = vec![]; + + for n in 0..=to { + debug!("Pushing {n} fibonacci", n = n); + sequence.push(nth_fibonacci(n)); + } + + sequence +} + +fn main() { + let layer = HierarchicalLayer::default() + .with_indent_lines(true) + .with_indent_amount(2) + .with_writer(std::io::stderr); + + let subscriber = Registry::default().with(layer); + tracing::subscriber::set_global_default(subscriber).unwrap(); + + let n = 5; + let sequence = fibonacci_seq(n); + info!("The first {} fibonacci numbers are {:?}", n, sequence); +} diff --git a/src/format.rs b/src/format.rs new file mode 100644 index 0000000..8a2c782 --- /dev/null +++ b/src/format.rs @@ -0,0 +1,196 @@ +use ansi_term::Color; +use std::{ + fmt::{self, Write as _}, + io, +}; +use tracing::{ + field::{Field, Visit}, + Level, +}; + +const LINE_VERT: &str = "│"; +const LINE_HORIZ: &str = "─"; +const LINE_BRANCH: &str = "├"; + +#[derive(Debug)] +pub struct Config { + pub ansi: bool, + pub indent_lines: bool, + pub indent_amount: usize, +} + +impl Config { + pub fn with_ansi(self, ansi: bool) -> Self { + Self { ansi, ..self } + } + + pub 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 Buffers { + pub current_buf: String, + pub indent_buf: String, +} + +impl Buffers { + pub fn new() -> Self { + Self { + current_buf: String::new(), + indent_buf: String::new(), + } + } + + pub fn flush_current_buf(&mut self, mut writer: impl io::Write) { + write!(writer, "{}", &self.current_buf).unwrap(); + self.current_buf.clear(); + } + + pub fn flush_indent_buf(&mut self) { + self.current_buf.push_str(&self.indent_buf); + self.indent_buf.clear(); + } + + pub fn indent_current(&mut self, indent: usize, config: &Config) { + indent_block( + &mut self.current_buf, + &mut self.indent_buf, + indent, + config.indent_amount, + config.indent_lines, + ); + self.current_buf.clear(); + } +} + +pub struct FmtEvent<'a> { + pub bufs: &'a mut Buffers, + pub comma: bool, +} + +impl<'a> Visit for FmtEvent<'a> { + fn record_debug(&mut self, field: &Field, value: &dyn fmt::Debug) { + let buf = &mut self.bufs.current_buf; + write!(buf, "{comma} ", comma = if self.comma { "," } else { "" },).unwrap(); + let name = field.name(); + if name == "message" { + write!(buf, "{:?}", value).unwrap(); + self.comma = true; + } else { + write!(buf, "{}={:?}", name, value).unwrap(); + self.comma = true; + } + } +} + +impl<'a> FmtEvent<'a> { + pub fn finish(&mut self, indent: usize, config: &Config) { + self.bufs.current_buf.push('\n'); + self.bufs.indent_current(indent, config); + self.bufs.flush_indent_buf(); + } +} + +pub struct ColorLevel<'a>(pub &'a Level); + +impl<'a> fmt::Display for ColorLevel<'a> { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + match *self.0 { + Level::TRACE => Color::Purple.bold().paint("TRACE"), + Level::DEBUG => Color::Blue.bold().paint("DEBUG"), + Level::INFO => Color::Green.bold().paint(" INFO"), + Level::WARN => Color::RGB(252, 234, 160).bold().paint(" WARN"), // orange + Level::ERROR => Color::Red.bold().paint("ERROR"), + } + .fmt(f) + } +} + +fn indent_block_with_lines(lines: &[&str], buf: &mut String, indent: usize, indent_amount: usize) { + let indent_spaces = indent * indent_amount; + if lines.is_empty() { + 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'); + } + } +} diff --git a/src/lib.rs b/src/lib.rs index 26c2ef0..b18aa9f 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -1,109 +1,31 @@ +pub(crate) mod format; + use ansi_term::{Color, Style}; use chrono::{DateTime, Local}; -use std::sync::Mutex; +use format::{Buffers, ColorLevel, Config, FmtEvent}; use std::{ fmt::{self, Write as _}, io, + sync::Mutex, }; - use tracing::{ field::{Field, Visit}, span::{Attributes, Id}, - Event, Level, Subscriber, + Event, Subscriber, }; use tracing_subscriber::{ + fmt::MakeWriter, layer::{Context, Layer}, 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, - bufs: Mutex, - config: Config, -} - -#[derive(Debug)] -struct Buffers { - pub current_buf: String, - pub indent_buf: String, -} - -impl Buffers { - fn new() -> Self { - Self { - current_buf: String::new(), - indent_buf: String::new(), - } - } - - fn flush_current_buf(&mut self, mut writer: impl io::Write) { - write!(writer, "{}", &self.current_buf).unwrap(); - self.current_buf.clear(); - } - - fn flush_indent_buf(&mut self) { - self.current_buf.push_str(&self.indent_buf); - self.indent_buf.clear(); - } - - fn indent_current(&mut self, indent: usize, config: &Config) { - indent_block( - &mut self.current_buf, - &mut self.indent_buf, - indent, - config.indent_amount, - config.indent_lines, - ); - self.current_buf.clear(); - } -} - -struct Data { +pub(crate) struct Data { start: DateTime, kvs: Vec<(&'static str, String)>, } -struct FmtEvent<'a> { - bufs: &'a mut Buffers, - comma: bool, -} - impl Data { - fn new(attrs: &tracing::span::Attributes<'_>) -> Self { + pub fn new(attrs: &tracing::span::Attributes<'_>) -> Self { let mut span = Self { start: Local::now(), kvs: Vec::new(), @@ -119,119 +41,34 @@ impl Visit for Data { } } -impl<'a> Visit for FmtEvent<'a> { - fn record_debug(&mut self, field: &Field, value: &dyn fmt::Debug) { - let buf = &mut self.bufs.current_buf; - write!(buf, "{comma} ", comma = if self.comma { "," } else { "" },).unwrap(); - let name = field.name(); - if name == "message" { - write!(buf, "{:?}", value).unwrap(); - self.comma = true; - } else { - write!(buf, "{}={:?}", name, value).unwrap(); - self.comma = true; - } - } -} - -impl<'a> FmtEvent<'a> { - fn finish(&mut self, indent: usize, config: &Config) { - self.bufs.current_buf.push('\n'); - self.bufs.indent_current(indent, config); - self.bufs.flush_indent_buf(); - } -} - -struct ColorLevel<'a>(&'a Level); - -impl<'a> fmt::Display for ColorLevel<'a> { - fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { - match *self.0 { - Level::TRACE => Color::Purple.bold().paint("TRACE"), - Level::DEBUG => Color::Blue.bold().paint("DEBUG"), - Level::INFO => Color::Green.bold().paint(" INFO"), - Level::WARN => Color::RGB(252, 234, 160).bold().paint(" WARN"), // orange - Level::ERROR => Color::Red.bold().paint("ERROR"), - } - .fmt(f) - } -} - -fn indent_block_with_lines(lines: &[&str], buf: &mut String, indent: usize, indent_amount: usize) { - let indent_spaces = indent * indent_amount; - 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'); - } +#[derive(Debug)] +pub struct HierarchicalLayer io::Stdout> +where + W: MakeWriter + 'static, +{ + make_writer: W, + bufs: Mutex, + config: Config, } -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 Default for HierarchicalLayer { + fn default() -> Self { + let ansi = atty::is(atty::Stream::Stdout); + let indent_amount = 2; + let config = Config { + ansi, + indent_amount, + ..Default::default() + }; + Self { + make_writer: io::stdout, + bufs: Mutex::new(Buffers::new()), + config, } } } -impl HierarchicalLayer { +impl HierarchicalLayer io::Stdout> { pub fn new(indent_amount: usize) -> Self { let ansi = atty::is(atty::Stream::Stdout); let config = Config { @@ -240,12 +77,17 @@ impl HierarchicalLayer { ..Default::default() }; Self { - stdout: io::stdout(), + make_writer: io::stdout, bufs: Mutex::new(Buffers::new()), config, } } +} +impl HierarchicalLayer +where + W: MakeWriter + 'static, +{ pub fn with_ansi(self, ansi: bool) -> Self { Self { config: self.config.with_ansi(ansi), @@ -253,6 +95,25 @@ impl HierarchicalLayer { } } + pub fn with_writer(self, make_writer: W2) -> HierarchicalLayer + where + W2: MakeWriter + 'static, + { + HierarchicalLayer { + make_writer, + config: self.config, + bufs: self.bufs, + } + } + + pub fn with_indent_amount(self, indent_amount: usize) -> Self { + let config = Config { + indent_amount, + ..self.config + }; + Self { config, ..self } + } + pub fn with_indent_lines(self, indent_lines: bool) -> Self { Self { config: self.config.with_indent_lines(indent_lines), @@ -290,9 +151,10 @@ impl HierarchicalLayer { } } -impl Layer for HierarchicalLayer +impl Layer for HierarchicalLayer where S: Subscriber + for<'span> LookupSpan<'span> + fmt::Debug, + W: MakeWriter + 'static, { fn new_span(&self, attrs: &Attributes, id: &Id, ctx: Context) { let data = Data::new(attrs); @@ -334,7 +196,8 @@ where bufs.indent_current(indent, &self.config); bufs.flush_indent_buf(); - bufs.flush_current_buf(self.stdout.lock()); + let writer = self.make_writer.make_writer(); + bufs.flush_current_buf(writer) } fn on_event(&self, event: &Event<'_>, ctx: Context) { @@ -392,7 +255,8 @@ where }; event.record(&mut visitor); visitor.finish(indent, &self.config); - bufs.flush_current_buf(self.stdout.lock()); + let writer = self.make_writer.make_writer(); + bufs.flush_current_buf(writer) } fn on_exit(&self, _id: &Id, _ctx: Context) {}