tracing-fancytree 0.1.1

#![doc = include_str!("../README.md")]

use {
	::core::fmt::{Debug, Write as _},
	::nu_ansi_term::{Color, Style},
	::std::{
		io::{self, IsTerminal, Write},
		sync::Mutex,
	},
	::tracing_core::{
		field::Visit,
		span::{Attributes, Id, Record},
		Event,
		Field,
		Level,
		Metadata,
		Subscriber,
	},
	::tracing_subscriber::{layer::Context, registry::LookupSpan, Layer},
};

fn span_continue_marker(sym: bool) -> &'static str {
	if sym {
		"│ "
	} else {
		"| "
	}
}
fn span_close_marker(sym: bool) -> &'static str {
	if sym {
		"╰·"
	} else {
		"\\-"
	}
}

fn event_marker(level: Level, sym: bool) -> (Style, &'static str) {
	(
		match level {
			Level::ERROR => Color::Red,
			Level::WARN => Color::Yellow,
			Level::INFO => Color::Blue,
			Level::DEBUG => Color::Magenta,
			Level::TRACE => Color::DarkGray,
		}
		.bold(),
		if sym {
			match level {
				Level::ERROR => " ",
				Level::WARN => " ",
				Level::INFO => " ",
				Level::DEBUG => " ",
				Level::TRACE => " ",
			}
		} else {
			match level {
				Level::ERROR => "[err] ",
				Level::WARN => "[wrn] ",
				Level::INFO => "[inf] ",
				Level::DEBUG => "[dbg] ",
				Level::TRACE => "[trc] ",
			}
		},
	)
}

/// print fancy tracing trees
pub struct FancyTree<O: Write + 'static> {
	inner: Mutex<Inner<O>>,
}

impl<O: Write + 'static> FancyTree<O> {
	/// setup a fancy tree printer
	///
	/// - `out` - the stream to write to
	/// - `ansi` - whether to print colors
	/// - `sym` - whether to print unicode symbols
	pub fn new(out: O, ansi: bool, sym: bool) -> Self {
		Self {
			inner: Mutex::new(Inner {
				ansi,
				sym,
				out,
				indent: Vec::new(),
			}),
		}
	}
}

impl Default for FancyTree<io::Stdout> {
	fn default() -> Self {
		let out = io::stdout();
		let ansi = out.is_terminal();
		Self::new(out, ansi, ansi)
	}
}

impl<S: Subscriber + for<'span> LookupSpan<'span, Data: Debug> + Debug, O: Write + 'static> Layer<S>
	for FancyTree<O>
{
	fn on_new_span(&self, span: &Attributes<'_>, _: &Id, _: Context<'_, S>) {
		let mut inner = self.inner.lock().unwrap();

		inner.line_prefix();

		let level = *span.metadata().level();
		let (marker_style, marker) = event_marker(level, inner.sym);
		if inner.ansi {
			let style = Color::Green.normal();
			write!(
				inner.out,
				"{}{marker}{}{} {}({}){}",
				marker_style.prefix(),
				span.metadata().name(),
				marker_style.suffix(),
				style.prefix(),
				span.metadata().target(),
				style.suffix(),
			)
			.unwrap();
		} else {
			write!(
				inner.out,
				"{marker}{} ({})",
				span.metadata().name(),
				span.metadata().target(),
			)
			.unwrap();
		}

		write_file(&mut *inner, span.metadata());

		writeln!(inner.out).unwrap();

		inner.indent.push(level);

		span.values().record(&mut FieldDisplayVisitor {
			inner: &mut inner,
			indent: "",
			ignore_message: false,
		});
	}

	fn on_record(&self, _: &Id, values: &Record<'_>, _: Context<'_, S>) {
		let mut inner = self.inner.lock().unwrap();
		values.record(&mut FieldDisplayVisitor {
			inner: &mut inner,
			indent: "",
			ignore_message: false,
		});
	}

	fn on_event(&self, event: &Event<'_>, _: Context<'_, S>) {
		let mut inner = self.inner.lock().unwrap();

		let mut return_finder = FieldVisitor {
			field: "return",
			output: None,
		};
		event.record(&mut return_finder);
		if let Some((_, false)) = return_finder.output {
			event.record(&mut FieldDisplayVisitor {
				inner: &mut inner,
				indent: "",
				ignore_message: false,
			});
			return;
		}

		let meta = event.metadata();
		let (style, sym) = event_marker(*meta.level(), inner.sym);

		let mut message_finder = FieldVisitor {
			field: "message",
			output: None,
		};
		event.record(&mut message_finder);
		if let Some((message, _)) = message_finder.output {
			inner.line_prefix();
			if inner.ansi {
				write!(inner.out, "{}{sym}", style.prefix()).unwrap();
			} else {
				write!(inner.out, "{sym}").unwrap();
			}

			if inner.ansi {
				write!(inner.out, "{}{message}{}", style.prefix(), style.suffix()).unwrap();
			} else {
				write!(inner.out, "{message}").unwrap();
			}
		} else {
			inner.line_prefix();
			if inner.ansi {
				write!(
					inner.out,
					"{}{sym}{}{}",
					style.prefix(),
					meta.name(),
					style.suffix()
				)
				.unwrap();
			} else {
				write!(inner.out, "{sym}{}", meta.name()).unwrap();
			}
		}

		write_file(&mut *inner, event.metadata());

		writeln!(inner.out).unwrap();

		event.record(&mut FieldDisplayVisitor {
			inner: &mut inner,
			indent: "  ",
			ignore_message: true,
		});
	}

	fn on_exit(&self, span: &Id, ctx: Context<'_, S>) {
		let mut inner = self.inner.lock().unwrap();

		let name = ctx.span(span).unwrap().name();
		let level = inner.indent.pop().unwrap();

		inner.line_prefix();

		let close_sym = span_close_marker(inner.sym);
		if inner.ansi {
			let style = event_marker(level, inner.sym).0;
			let name_style = Color::DarkGray.normal().bold();
			writeln!(
				inner.out,
				"{}{close_sym}{}{name}{}",
				style.prefix(),
				style.infix(name_style),
				name_style.suffix()
			)
			.unwrap();
		} else {
			writeln!(inner.out, "{close_sym} {name}").unwrap();
		}
	}
}

fn write_file<O: Write + 'static>(inner: &mut Inner<O>, metadata: &'static Metadata<'static>) {
	if metadata.file().is_some() || metadata.line().is_some() {
		if inner.ansi {
			let style = Color::DarkGray.normal();
			write!(
				inner.out,
				" {}{}:{}{}",
				style.prefix(),
				metadata.file().unwrap_or_default(),
				metadata
					.line()
					.map_or_else(String::new, |line| line.to_string()),
				style.suffix(),
			)
			.unwrap();
		} else {
			write!(
				inner.out,
				" {}:{}",
				metadata.file().unwrap_or_default(),
				metadata
					.line()
					.map_or_else(String::new, |line| line.to_string()),
			)
			.unwrap();
		}
	}
}

struct Inner<O: Write + 'static> {
	ansi: bool,
	sym: bool,
	out: O,
	indent: Vec<Level>,
}

impl<O: Write + 'static> Inner<O> {
	fn line_prefix(&mut self) {
		let continue_sym = span_continue_marker(self.sym);
		if self.ansi {
			let mut last_color = Style::new();
			for level in &self.indent {
				let style = event_marker(*level, self.sym).0;
				write!(self.out, "{}{continue_sym}", last_color.infix(style)).unwrap();
				last_color = style;
			}
			write!(self.out, "{}", last_color.suffix()).unwrap();
		} else {
			for _ in &self.indent {
				write!(self.out, "{continue_sym}").unwrap();
			}
		}
	}
}

struct FieldDisplayVisitor<'inner, O: Write + 'static> {
	inner: &'inner mut Inner<O>,
	indent: &'static str,
	ignore_message: bool,
}

impl<O: Write + 'static> FieldDisplayVisitor<'_, O> {
	fn field(&mut self, name: &str, style: Style, content: &str) {
		if name == "message" && self.ignore_message {
			return;
		}

		for (i, line) in content.lines().enumerate() {
			self.inner.line_prefix();

			write!(self.inner.out, "{}", self.indent).unwrap();

			if i == 0 {
				if self.inner.ansi {
					let name_style = Color::LightGreen.normal().bold();
					write!(
						self.inner.out,
						"{}{name}{}: ",
						name_style.prefix(),
						name_style.suffix()
					)
					.unwrap();
				} else {
					write!(self.inner.out, "{name}: ").unwrap();
				}
			} else {
				write!(self.inner.out, "{}  ", " ".repeat(name.len())).unwrap();
			}

			if self.inner.ansi {
				write!(self.inner.out, "{}{line}{}", style.prefix(), style.suffix()).unwrap();
			} else {
				write!(self.inner.out, "{line}").unwrap();
			}

			writeln!(self.inner.out).unwrap();
		}
	}
}

impl<O: Write + 'static> Visit for FieldDisplayVisitor<'_, O> {
	fn record_f64(&mut self, field: &Field, value: f64) {
		self.field(field.name(), Color::Yellow.normal(), &format!("{value}"));
	}

	fn record_i64(&mut self, field: &Field, value: i64) {
		self.field(field.name(), Color::Yellow.normal(), &format!("{value}"));
	}

	fn record_u64(&mut self, field: &Field, value: u64) {
		self.field(field.name(), Color::Yellow.normal(), &format!("{value}"));
	}

	fn record_i128(&mut self, field: &Field, value: i128) {
		self.field(field.name(), Color::Yellow.normal(), &format!("{value}"));
	}

	fn record_u128(&mut self, field: &Field, value: u128) {
		self.field(field.name(), Color::Yellow.normal(), &format!("{value}"));
	}

	fn record_bool(&mut self, field: &Field, value: bool) {
		self.field(
			field.name(),
			Color::LightGreen.normal(),
			&format!("{value}"),
		);
	}

	fn record_str(&mut self, field: &Field, value: &str) {
		self.field(field.name(), Color::Blue.normal(), value);
	}

	fn record_bytes(&mut self, field: &Field, value: &[u8]) {
		let mut hex = String::with_capacity((value.len() * 3).saturating_sub(1));

		for (i, byte) in value.iter().enumerate() {
			if i != 0 {
				hex.push(' ');
			}
			write!(hex, "{byte:x}").unwrap();
		}

		self.field(field.name(), Color::Purple.normal(), &hex);
	}

	fn record_error(&mut self, field: &Field, value: &(dyn ::std::error::Error + 'static)) {
		self.field(field.name(), Color::Red.normal(), &format!("{value}"));
	}

	fn record_debug(&mut self, field: &Field, value: &dyn Debug) {
		self.field(field.name(), Color::White.normal(), &format!("{value:#?}"));
	}
}

struct FieldVisitor {
	field: &'static str,
	output: Option<(String, bool)>,
}

impl Visit for FieldVisitor {
	fn record_str(&mut self, field: &Field, value: &str) {
		if let Some((_, any_others)) = &mut self.output {
			*any_others = true;
		}

		if field.name() == self.field {
			self.output = Some((value.to_string(), false));
		}
	}

	fn record_debug(&mut self, field: &Field, value: &dyn Debug) {
		if let Some((_, any_others)) = &mut self.output {
			*any_others = true;
		}

		if field.name() == self.field {
			self.output = Some((format!("{value:#?}"), false));
		}
	}
}