use crate::{OutputFormat, console_fields::ConsoleFields, console_format::ConsoleFormat};
use std::{
io::{self, IsTerminal, Write},
sync::Mutex,
time::Instant,
};
use tracing::{Subscriber, subscriber::SetGlobalDefaultError};
use tracing_error::ErrorLayer;
use tracing_subscriber::{
EnvFilter, Layer, Registry, filter::FromEnvError, layer::SubscriberExt, registry::LookupSpan,
};
pub struct Console<W = io::Stderr> {
writer: W,
format: OutputFormat,
terminal: bool,
ansi: Option<bool>,
filter: EnvFilter,
}
impl Default for Console {
fn default() -> Self {
let writer = io::stderr();
let terminal = writer.is_terminal();
Self {
writer,
format: OutputFormat::Auto,
terminal,
ansi: None,
filter: EnvFilter::new("info"),
}
}
}
impl<W> Console<W> {
pub fn with_format(mut self, format: OutputFormat) -> Self {
self.format = format;
self
}
pub fn with_ansi(mut self, enabled: bool) -> Self {
self.ansi = Some(enabled);
self
}
pub fn with_writer<T>(self, writer: T) -> Console<T> {
Console {
writer,
format: self.format,
terminal: false,
ansi: self.ansi,
filter: self.filter,
}
}
pub fn with_filter(mut self, filter: EnvFilter) -> Self {
self.filter = filter;
self
}
pub fn with_env_filter(mut self) -> Result<Self, FromEnvError> {
self.filter = EnvFilter::builder()
.with_default_directive(tracing::level_filters::LevelFilter::INFO.into())
.from_env()?;
Ok(self)
}
}
impl<W: Write + Send + 'static> Console<W> {
pub fn layer<S>(self) -> impl Layer<S>
where
S: Subscriber + for<'lookup> LookupSpan<'lookup>,
{
let filter = self.filter.clone();
self.output_layer().with_filter(filter)
}
fn output_layer<S>(self) -> impl Layer<S>
where
S: Subscriber + for<'lookup> LookupSpan<'lookup>,
{
let format = self.format.resolve(self.terminal);
let ansi = self.ansi.unwrap_or(self.terminal) && format == OutputFormat::Text;
let formatter = ConsoleFormat {
format,
started: Instant::now(),
ansi,
};
let output = tracing_subscriber::fmt::layer()
.fmt_fields(ConsoleFields)
.event_format(formatter)
.with_ansi(ansi)
.with_writer(Mutex::new(RecordWriter(self.writer)));
ErrorLayer::new(ConsoleFields).and_then(output)
}
pub fn subscriber(self) -> impl Subscriber + Send + Sync {
let filter = self.filter.clone();
Registry::default().with(filter).with(self.output_layer())
}
#[track_caller]
pub fn init(self) {
self.try_init().expect("could not initialize the console");
}
pub fn try_init(self) -> Result<(), SetGlobalDefaultError> {
tracing::subscriber::set_global_default(self.subscriber())
}
}
struct RecordWriter<W>(W);
impl<W: Write> Write for RecordWriter<W> {
fn write(&mut self, bytes: &[u8]) -> io::Result<usize> {
let written = self.0.write(bytes)?;
self.0.flush()?;
Ok(written)
}
fn flush(&mut self) -> io::Result<()> {
self.0.flush()
}
}
#[cfg(test)]
mod tests;