use std::any::TypeId;
use std::marker::PhantomData;
use std::sync::Mutex;
use indicatif::MultiProgress;
use indicatif::ProgressBar;
pub use indicatif::style;
use indicatif::style::ProgressStyle;
use indicatif::style::ProgressTracker;
use tracing_core::Subscriber;
use tracing_core::span;
use tracing_subscriber::fmt::FormatFields;
use tracing_subscriber::fmt::FormattedFields;
use tracing_subscriber::fmt::format::DefaultFields;
use tracing_subscriber::layer;
use tracing_subscriber::registry::LookupSpan;
pub mod filter;
mod pb_manager;
pub mod span_ext;
pub mod util;
pub mod writer;
use pb_manager::ProgressBarManager;
pub use pb_manager::TickSettings;
#[doc(inline)]
pub use writer::IndicatifWriter;
#[derive(Clone)]
struct IndicatifProgressKey {
message: String,
}
impl ProgressTracker for IndicatifProgressKey {
fn clone_box(&self) -> Box<dyn ProgressTracker> {
Box::new(self.clone())
}
fn tick(&mut self, _: &indicatif::ProgressState, _: std::time::Instant) {}
fn reset(&mut self, _: &indicatif::ProgressState, _: std::time::Instant) {}
fn write(&self, _: &indicatif::ProgressState, w: &mut dyn std::fmt::Write) {
let _ = w.write_str(&self.message);
}
}
#[allow(clippy::type_complexity)]
pub(crate) struct WithContext(
fn(&tracing::Dispatch, &span::Id, f: &mut dyn FnMut(&mut IndicatifSpanContext)),
);
#[allow(clippy::type_complexity)]
pub(crate) struct WithStderrWriter(
fn(&tracing::Dispatch, f: &mut dyn FnMut(IndicatifWriter<writer::Stderr>)),
);
#[allow(clippy::type_complexity)]
pub(crate) struct WithStdoutWriter(
fn(&tracing::Dispatch, f: &mut dyn FnMut(IndicatifWriter<writer::Stdout>)),
);
#[allow(clippy::type_complexity)]
pub(crate) struct WithMultiProgress(fn(&tracing::Dispatch, f: &mut dyn FnMut(MultiProgress)));
impl WithContext {
pub(crate) fn with_context(
&self,
dispatch: &tracing::Dispatch,
id: &span::Id,
mut f: impl FnMut(&mut IndicatifSpanContext),
) {
(self.0)(dispatch, id, &mut f)
}
}
impl WithStderrWriter {
pub(crate) fn with_context(
&self,
dispatch: &tracing::Dispatch,
mut f: impl FnMut(IndicatifWriter<writer::Stderr>),
) {
(self.0)(dispatch, &mut f)
}
}
impl WithStdoutWriter {
pub(crate) fn with_context(
&self,
dispatch: &tracing::Dispatch,
mut f: impl FnMut(IndicatifWriter<writer::Stdout>),
) {
(self.0)(dispatch, &mut f)
}
}
impl WithMultiProgress {
pub(crate) fn with_context(
&self,
dispatch: &tracing::Dispatch,
mut f: impl FnMut(MultiProgress),
) {
(self.0)(dispatch, &mut f)
}
}
#[derive(Default)]
struct ProgressBarInitSettings {
style: Option<ProgressStyle>,
len: Option<u64>,
pos: Option<u64>,
message: Option<String>,
}
struct IndicatifSpanContext {
progress_bar: Option<ProgressBar>,
pb_init_settings: ProgressBarInitSettings,
parent_progress_bar: Option<ProgressBar>,
parent_span: Option<span::Id>,
span_fields_formatted: Option<String>,
span_name: String,
span_child_prefix: String,
level: u16,
finish_message: Option<String>,
}
impl IndicatifSpanContext {
fn add_keys_to_style(&self, style: ProgressStyle) -> ProgressStyle {
style
.with_key(
"span_name",
IndicatifProgressKey {
message: self.span_name.clone(),
},
)
.with_key(
"span_fields",
IndicatifProgressKey {
message: self.span_fields_formatted.to_owned().unwrap_or_default(),
},
)
.with_key(
"span_child_prefix",
IndicatifProgressKey {
message: self.span_child_prefix.clone(),
},
)
}
fn make_progress_bar(&mut self, default_style: &ProgressStyle) {
if self.progress_bar.is_none() {
let pb = ProgressBar::hidden().with_style(
self.pb_init_settings
.style
.take()
.unwrap_or_else(|| self.add_keys_to_style(default_style.clone())),
);
if let Some(len) = self.pb_init_settings.len.take() {
pb.set_length(len);
}
if let Some(msg) = self.pb_init_settings.message.take() {
pb.set_message(msg);
}
if let Some(pos) = self.pb_init_settings.pos.take() {
pb.set_position(pos);
}
self.progress_bar = Some(pb);
}
}
fn set_progress_bar_style(&mut self, style: ProgressStyle) {
if let Some(ref pb) = self.progress_bar {
pb.set_style(self.add_keys_to_style(style));
} else {
self.pb_init_settings.style = Some(self.add_keys_to_style(style));
}
}
fn set_progress_bar_length(&mut self, len: u64) {
if let Some(ref pb) = self.progress_bar {
pb.set_length(len);
} else {
self.pb_init_settings.len = Some(len);
}
}
fn set_progress_bar_position(&mut self, pos: u64) {
if let Some(ref pb) = self.progress_bar {
pb.set_position(pos);
} else {
self.pb_init_settings.pos = Some(pos);
}
}
fn set_progress_bar_message(&mut self, msg: String) {
if let Some(ref pb) = self.progress_bar {
pb.set_message(msg);
} else {
self.pb_init_settings.message = Some(msg);
}
}
fn inc_progress_bar_position(&mut self, pos: u64) {
if let Some(ref pb) = self.progress_bar {
pb.inc(pos);
} else if let Some(ref mut pb_pos) = self.pb_init_settings.pos {
*pb_pos += pos;
} else {
self.pb_init_settings.pos = Some(pos);
}
}
fn inc_progress_bar_length(&mut self, len: u64) {
if let Some(ref pb) = self.progress_bar {
pb.inc_length(len);
} else if let Some(ref mut pb_len) = self.pb_init_settings.len {
*pb_len += len;
}
}
fn progress_bar_tick(&mut self) {
if let Some(ref pb) = self.progress_bar {
pb.tick()
}
}
fn reset_progress_bar(&mut self) {
if let Some(ref pb) = self.progress_bar {
pb.reset();
}
}
fn reset_progress_bar_elapsed(&mut self) {
if let Some(ref pb) = self.progress_bar {
pb.reset_elapsed();
}
}
fn reset_progress_bar_eta(&mut self) {
if let Some(ref pb) = self.progress_bar {
pb.reset_eta();
}
}
fn set_progress_bar_finish_message(&mut self, msg: String) {
self.finish_message = Some(msg);
}
}
pub struct IndicatifLayer<S, F = DefaultFields> {
pb_manager: Mutex<ProgressBarManager>,
mp: MultiProgress,
span_field_formatter: F,
progress_style: ProgressStyle,
span_child_prefix_indent: &'static str,
span_child_prefix_symbol: &'static str,
get_context: WithContext,
get_stderr_writer_context: WithStderrWriter,
get_stdout_writer_context: WithStdoutWriter,
get_multi_progress_context: WithMultiProgress,
inner: PhantomData<S>,
}
impl<S> IndicatifLayer<S>
where
S: Subscriber + for<'a> LookupSpan<'a>,
{
pub fn new() -> Self {
Self::default()
}
}
impl<S> Default for IndicatifLayer<S>
where
S: Subscriber + for<'a> LookupSpan<'a>,
{
fn default() -> Self {
let pb_manager = ProgressBarManager::new(
7,
Some(
ProgressStyle::with_template(
"...and {pending_progress_bars} more not shown above.",
)
.expect("valid template"),
),
TickSettings::default(),
);
let mp = pb_manager.mp.clone();
Self {
pb_manager: Mutex::new(pb_manager),
mp,
span_field_formatter: DefaultFields::new(),
progress_style: ProgressStyle::with_template(
"{span_child_prefix}{spinner} {span_name}{{{span_fields}}}",
)
.expect("valid template"),
span_child_prefix_indent: " ",
span_child_prefix_symbol: "↳ ",
get_context: WithContext(Self::get_context),
get_stderr_writer_context: WithStderrWriter(Self::get_stderr_writer_context),
get_stdout_writer_context: WithStdoutWriter(Self::get_stdout_writer_context),
get_multi_progress_context: WithMultiProgress(Self::get_multi_progress_context),
inner: PhantomData,
}
}
}
impl<S, F> IndicatifLayer<S, F>
where
S: Subscriber + for<'a> LookupSpan<'a>,
{
#[deprecated(since = "0.2.3", note = "use get_stderr_writer() instead")]
pub fn get_fmt_writer(&self) -> IndicatifWriter<writer::Stderr> {
self.get_stderr_writer()
}
pub fn get_stderr_writer(&self) -> IndicatifWriter<writer::Stderr> {
IndicatifWriter::new(self.mp.clone())
}
pub fn get_stdout_writer(&self) -> IndicatifWriter<writer::Stdout> {
IndicatifWriter::new(self.mp.clone())
}
pub fn with_span_field_formatter<F2>(self, formatter: F2) -> IndicatifLayer<S, F2>
where
F2: for<'writer> FormatFields<'writer> + 'static,
{
IndicatifLayer {
pb_manager: self.pb_manager,
mp: self.mp,
span_field_formatter: formatter,
progress_style: self.progress_style,
span_child_prefix_indent: self.span_child_prefix_indent,
span_child_prefix_symbol: self.span_child_prefix_symbol,
get_context: WithContext(IndicatifLayer::<S, F2>::get_context),
get_stderr_writer_context: WithStderrWriter(
IndicatifLayer::<S, F2>::get_stderr_writer_context,
),
get_stdout_writer_context: WithStdoutWriter(
IndicatifLayer::<S, F2>::get_stdout_writer_context,
),
get_multi_progress_context: WithMultiProgress(
IndicatifLayer::<S, F2>::get_multi_progress_context,
),
inner: self.inner,
}
}
pub fn with_progress_style(mut self, style: ProgressStyle) -> Self {
self.progress_style = style;
self
}
pub fn with_span_child_prefix_indent(mut self, indent: &'static str) -> Self {
self.span_child_prefix_indent = indent;
self
}
pub fn with_span_child_prefix_symbol(mut self, symbol: &'static str) -> Self {
self.span_child_prefix_symbol = symbol;
self
}
pub fn with_max_progress_bars(
mut self,
max_progress_bars: u64,
footer_style: Option<ProgressStyle>,
) -> Self {
if let Ok(pb_manager) = self.pb_manager.get_mut() {
pb_manager.set_max_progress_bars(max_progress_bars, footer_style);
}
self
}
pub fn with_tick_settings(mut self, tick_settings: TickSettings) -> Self {
if let Ok(pb_manager) = self.pb_manager.get_mut() {
pb_manager.set_tick_settings(tick_settings);
}
self
}
}
impl<S, F> IndicatifLayer<S, F>
where
S: Subscriber + for<'a> LookupSpan<'a>,
F: for<'writer> FormatFields<'writer> + 'static,
{
fn get_context(
dispatch: &tracing::Dispatch,
id: &span::Id,
f: &mut dyn FnMut(&mut IndicatifSpanContext),
) {
let subscriber = dispatch
.downcast_ref::<S>()
.expect("subscriber should downcast to expected type; this is a bug!");
let span = subscriber
.span(id)
.expect("Span not found in context, this is a bug");
let mut ext = span.extensions_mut();
if let Some(indicatif_ctx) = ext.get_mut::<IndicatifSpanContext>() {
f(indicatif_ctx);
}
}
fn get_stderr_writer_context(
dispatch: &tracing::Dispatch,
f: &mut dyn FnMut(IndicatifWriter<writer::Stderr>),
) {
let layer = dispatch
.downcast_ref::<IndicatifLayer<S, F>>()
.expect("subscriber should downcast to expected type; this is a bug!");
f(layer.get_stderr_writer())
}
fn get_stdout_writer_context(
dispatch: &tracing::Dispatch,
f: &mut dyn FnMut(IndicatifWriter<writer::Stdout>),
) {
let layer = dispatch
.downcast_ref::<IndicatifLayer<S, F>>()
.expect("subscriber should downcast to expected type; this is a bug!");
f(layer.get_stdout_writer())
}
fn get_multi_progress_context(dispatch: &tracing::Dispatch, f: &mut dyn FnMut(MultiProgress)) {
let layer = dispatch
.downcast_ref::<IndicatifLayer<S, F>>()
.expect("subscriber should downcast to expected type; this is a bug!");
f(layer.mp.clone())
}
fn handle_on_enter(
&self,
pb_manager: &mut ProgressBarManager,
id: &span::Id,
ctx: &layer::Context<'_, S>,
) -> Option<ProgressBar> {
let span = ctx
.span(id)
.expect("Span not found in context, this is a bug");
let mut ext = span.extensions_mut();
if let Some(indicatif_ctx) = ext.get_mut::<IndicatifSpanContext>() {
if indicatif_ctx.progress_bar.is_none() {
indicatif_ctx.make_progress_bar(&self.progress_style);
if let Some(ref parent_span_with_pb) = indicatif_ctx.parent_span {
let parent_pb = self.handle_on_enter(pb_manager, parent_span_with_pb, ctx);
indicatif_ctx.parent_progress_bar = parent_pb;
}
pb_manager.show_progress_bar(indicatif_ctx, id);
}
return indicatif_ctx.progress_bar.to_owned();
}
None
}
}
impl<S, F> layer::Layer<S> for IndicatifLayer<S, F>
where
S: Subscriber + for<'a> LookupSpan<'a>,
F: for<'writer> FormatFields<'writer> + 'static,
{
fn on_new_span(&self, attrs: &span::Attributes<'_>, id: &span::Id, ctx: layer::Context<'_, S>) {
let span = ctx
.span(id)
.expect("Span not found in context, this is a bug");
let mut ext = span.extensions_mut();
let mut fields = FormattedFields::<F>::new(String::new());
let _ = self
.span_field_formatter
.format_fields(fields.as_writer(), attrs);
let parent_span = ctx.span_scope(id).and_then(|scope| {
scope.skip(1).find(|span| {
let ext = span.extensions();
ext.get::<IndicatifSpanContext>().is_some()
})
});
let parent_span_id = parent_span.as_ref().map(|span| span.id());
let parent_span_ext = parent_span.as_ref().map(|span| span.extensions());
let parent_indicatif_ctx = parent_span_ext.as_ref().map(|ext| {
ext.get::<IndicatifSpanContext>()
.expect("validated it exists prior")
});
let (span_child_prefix, level) = match parent_indicatif_ctx {
Some(v) => {
let level = v.level + 1;
(
format!(
"{}{}",
self.span_child_prefix_indent.repeat(level.into()),
self.span_child_prefix_symbol
),
level,
)
}
None => (String::new(), 0),
};
ext.insert(IndicatifSpanContext {
progress_bar: None,
pb_init_settings: ProgressBarInitSettings::default(),
parent_progress_bar: None,
parent_span: parent_span_id,
span_fields_formatted: Some(fields.fields),
span_name: span.name().to_string(),
span_child_prefix,
level,
finish_message: None,
});
}
fn on_enter(&self, id: &span::Id, ctx: layer::Context<'_, S>) {
if let Ok(mut pb_manager_lock) = self.pb_manager.lock() {
self.handle_on_enter(&mut pb_manager_lock, id, &ctx);
}
}
fn on_close(&self, id: span::Id, ctx: layer::Context<'_, S>) {
if let Ok(mut pb_manager_lock) = self.pb_manager.lock() {
let span = ctx
.span(&id)
.expect("Span not found in context, this is a bug");
let mut ext = span.extensions_mut();
if let Some(indicatif_ctx) = ext.get_mut::<IndicatifSpanContext>() {
pb_manager_lock.finish_progress_bar(indicatif_ctx, &ctx);
}
}
}
unsafe fn downcast_raw(&self, id: TypeId) -> Option<*const ()> {
match id {
id if id == TypeId::of::<Self>() => Some(self as *const _ as *const ()),
id if id == TypeId::of::<WithContext>() => {
Some(&self.get_context as *const _ as *const ())
}
id if id == TypeId::of::<WithStderrWriter>() => {
Some(&self.get_stderr_writer_context as *const _ as *const ())
}
id if id == TypeId::of::<WithStdoutWriter>() => {
Some(&self.get_stdout_writer_context as *const _ as *const ())
}
id if id == TypeId::of::<WithMultiProgress>() => {
Some(&self.get_multi_progress_context as *const _ as *const ())
}
_ => None,
}
}
}
pub fn suspend_tracing_indicatif<F: FnOnce() -> R, R>(f: F) -> R {
let mut mp: Option<MultiProgress> = None;
tracing::dispatcher::get_default(|dispatch| {
if let Some(ctx) = dispatch.downcast_ref::<WithMultiProgress>() {
ctx.with_context(dispatch, |fetched_mp| {
mp = Some(fetched_mp);
})
}
});
if let Some(mp) = mp {
mp.suspend(f)
} else {
f()
}
}
#[macro_export]
macro_rules! indicatif_println {
($($arg:tt)*) => {
{
use std::io::Write;
if let Some(mut writer) = $crate::writer::get_indicatif_stdout_writer() {
writeln!(writer, $($arg)*).unwrap();
} else {
#[allow(clippy::explicit_write)]
writeln!(std::io::stdout(), $($arg)*).unwrap();
}
}
}
}
#[macro_export]
macro_rules! indicatif_eprintln {
($($arg:tt)*) => {
{
use std::io::Write;
if let Some(mut writer) = $crate::writer::get_indicatif_stderr_writer() {
writeln!(writer, $($arg)*).unwrap();
} else {
#[allow(clippy::explicit_write)]
writeln!(std::io::stderr(), $($arg)*).unwrap();
}
}
}
}
#[cfg(test)]
mod tests;