1use std::any::TypeId;
40use std::marker::PhantomData;
41use std::sync::Mutex;
42use std::time::Duration;
43
44use indicatif::MultiProgress;
45use indicatif::ProgressBar;
46pub use indicatif::style;
48use indicatif::style::ProgressStyle;
49use indicatif::style::ProgressTracker;
50use tracing_core::Subscriber;
51use tracing_core::span;
52use tracing_subscriber::fmt::FormatFields;
53use tracing_subscriber::fmt::FormattedFields;
54use tracing_subscriber::fmt::format::DefaultFields;
55use tracing_subscriber::layer;
56use tracing_subscriber::registry::LookupSpan;
57
58pub mod filter;
59mod pb_manager;
60pub mod span_ext;
61pub mod util;
62pub mod writer;
63
64use pb_manager::ProgressBarManager;
65pub use pb_manager::TickSettings;
66#[doc(inline)]
67pub use writer::IndicatifWriter;
68
69#[derive(Clone)]
70struct IndicatifProgressKey {
71 message: String,
72}
73
74impl ProgressTracker for IndicatifProgressKey {
75 fn clone_box(&self) -> Box<dyn ProgressTracker> {
76 Box::new(self.clone())
77 }
78
79 fn tick(&mut self, _: &indicatif::ProgressState, _: std::time::Instant) {}
80
81 fn reset(&mut self, _: &indicatif::ProgressState, _: std::time::Instant) {}
82
83 fn write(&self, _: &indicatif::ProgressState, w: &mut dyn std::fmt::Write) {
84 let _ = w.write_str(&self.message);
85 }
86}
87
88#[allow(clippy::type_complexity)]
111pub(crate) struct WithContext(
112 fn(&tracing::Dispatch, &span::Id, f: &mut dyn FnMut(&mut IndicatifSpanContext)),
113);
114
115#[allow(clippy::type_complexity)]
116pub(crate) struct WithStderrWriter(
117 fn(&tracing::Dispatch, f: &mut dyn FnMut(IndicatifWriter<writer::Stderr>)),
118);
119
120#[allow(clippy::type_complexity)]
121pub(crate) struct WithStdoutWriter(
122 fn(&tracing::Dispatch, f: &mut dyn FnMut(IndicatifWriter<writer::Stdout>)),
123);
124
125#[allow(clippy::type_complexity)]
126pub(crate) struct WithMultiProgress(fn(&tracing::Dispatch, f: &mut dyn FnMut(MultiProgress)));
127
128impl WithContext {
129 pub(crate) fn with_context(
130 &self,
131 dispatch: &tracing::Dispatch,
132 id: &span::Id,
133 mut f: impl FnMut(&mut IndicatifSpanContext),
134 ) {
135 (self.0)(dispatch, id, &mut f)
136 }
137}
138
139impl WithStderrWriter {
140 pub(crate) fn with_context(
141 &self,
142 dispatch: &tracing::Dispatch,
143 mut f: impl FnMut(IndicatifWriter<writer::Stderr>),
144 ) {
145 (self.0)(dispatch, &mut f)
146 }
147}
148
149impl WithStdoutWriter {
150 pub(crate) fn with_context(
151 &self,
152 dispatch: &tracing::Dispatch,
153 mut f: impl FnMut(IndicatifWriter<writer::Stdout>),
154 ) {
155 (self.0)(dispatch, &mut f)
156 }
157}
158
159impl WithMultiProgress {
160 pub(crate) fn with_context(
161 &self,
162 dispatch: &tracing::Dispatch,
163 mut f: impl FnMut(MultiProgress),
164 ) {
165 (self.0)(dispatch, &mut f)
166 }
167}
168
169#[derive(Default)]
170struct ProgressBarInitSettings {
171 style: Option<ProgressStyle>,
172 len: Option<u64>,
173 pos: Option<u64>,
174 message: Option<String>,
175}
176
177struct IndicatifSpanContext {
178 progress_bar: Option<ProgressBar>,
183 pb_init_settings: ProgressBarInitSettings,
185 parent_progress_bar: Option<ProgressBar>,
190 parent_span: Option<span::Id>,
192 span_fields_formatted: Option<String>,
194 span_name: String,
195 span_child_prefix: String,
196 level: u16,
199 finish_message: Option<String>,
201}
202
203impl IndicatifSpanContext {
204 fn add_keys_to_style(&self, style: ProgressStyle) -> ProgressStyle {
205 style
206 .with_key(
207 "span_name",
208 IndicatifProgressKey {
209 message: self.span_name.clone(),
210 },
211 )
212 .with_key(
213 "span_fields",
214 IndicatifProgressKey {
215 message: self.span_fields_formatted.to_owned().unwrap_or_default(),
216 },
217 )
218 .with_key(
219 "span_child_prefix",
220 IndicatifProgressKey {
221 message: self.span_child_prefix.clone(),
222 },
223 )
224 }
225
226 fn make_progress_bar(&mut self, default_style: &ProgressStyle) {
227 if self.progress_bar.is_none() {
228 let pb = ProgressBar::hidden().with_style(
229 self.pb_init_settings
230 .style
231 .take()
232 .unwrap_or_else(|| self.add_keys_to_style(default_style.clone())),
233 );
234
235 if let Some(len) = self.pb_init_settings.len.take() {
236 pb.set_length(len);
237 }
238
239 if let Some(msg) = self.pb_init_settings.message.take() {
240 pb.set_message(msg);
241 }
242
243 if let Some(pos) = self.pb_init_settings.pos.take() {
244 pb.set_position(pos);
245 }
246
247 self.progress_bar = Some(pb);
248 }
249 }
250
251 fn set_progress_bar_style(&mut self, style: ProgressStyle) {
252 if let Some(ref pb) = self.progress_bar {
253 pb.set_style(self.add_keys_to_style(style));
254 } else {
255 self.pb_init_settings.style = Some(self.add_keys_to_style(style));
256 }
257 }
258
259 fn set_progress_bar_length(&mut self, len: u64) {
260 if let Some(ref pb) = self.progress_bar {
261 pb.set_length(len);
262 } else {
263 self.pb_init_settings.len = Some(len);
264 }
265 }
266
267 fn set_progress_bar_position(&mut self, pos: u64) {
268 if let Some(ref pb) = self.progress_bar {
269 pb.set_position(pos);
270 } else {
271 self.pb_init_settings.pos = Some(pos);
272 }
273 }
274
275 fn set_progress_bar_message(&mut self, msg: String) {
276 if let Some(ref pb) = self.progress_bar {
277 pb.set_message(msg);
278 } else {
279 self.pb_init_settings.message = Some(msg);
280 }
281 }
282
283 fn inc_progress_bar_position(&mut self, pos: u64) {
284 if let Some(ref pb) = self.progress_bar {
285 pb.inc(pos);
286 } else if let Some(ref mut pb_pos) = self.pb_init_settings.pos {
287 *pb_pos += pos;
288 } else {
289 self.pb_init_settings.pos = Some(pos);
291 }
292 }
293
294 fn inc_progress_bar_length(&mut self, len: u64) {
295 if let Some(ref pb) = self.progress_bar {
296 pb.inc_length(len);
297 } else if let Some(ref mut pb_len) = self.pb_init_settings.len {
298 *pb_len += len;
299 }
300 }
301
302 fn progress_bar_tick(&mut self) {
303 if let Some(ref pb) = self.progress_bar {
304 pb.tick()
305 }
306 }
307
308 fn reset_progress_bar(&mut self) {
309 if let Some(ref pb) = self.progress_bar {
310 pb.reset();
311 }
312 }
313
314 fn reset_progress_bar_elapsed(&mut self) {
315 if let Some(ref pb) = self.progress_bar {
316 pb.reset_elapsed();
317 }
318 }
319
320 fn reset_progress_bar_eta(&mut self) {
321 if let Some(ref pb) = self.progress_bar {
322 pb.reset_eta();
323 }
324 }
325
326 fn set_progress_bar_finish_message(&mut self, msg: String) {
327 self.finish_message = Some(msg);
328 }
329
330 fn eta(&self) -> Duration {
331 if let Some(ref pb) = self.progress_bar {
332 pb.eta()
333 } else {
334 Duration::new(0, 0)
335 }
336 }
337
338 fn elapsed(&self) -> Duration {
339 if let Some(ref pb) = self.progress_bar {
340 pb.elapsed()
341 } else {
342 Duration::new(0, 0)
343 }
344 }
345}
346
347pub struct IndicatifLayer<S, F = DefaultFields> {
365 pb_manager: Mutex<ProgressBarManager>,
366 mp: MultiProgress,
369 span_field_formatter: F,
370 progress_style: ProgressStyle,
371 span_child_prefix_indent: &'static str,
372 span_child_prefix_symbol: &'static str,
373 get_context: WithContext,
374 get_stderr_writer_context: WithStderrWriter,
375 get_stdout_writer_context: WithStdoutWriter,
376 get_multi_progress_context: WithMultiProgress,
377 inner: PhantomData<S>,
378}
379
380impl<S> IndicatifLayer<S>
381where
382 S: Subscriber + for<'a> LookupSpan<'a>,
383{
384 pub fn new() -> Self {
399 Self::default()
400 }
401}
402
403impl<S> Default for IndicatifLayer<S>
404where
405 S: Subscriber + for<'a> LookupSpan<'a>,
406{
407 fn default() -> Self {
408 let pb_manager = ProgressBarManager::new(
409 7,
410 Some(
411 ProgressStyle::with_template(
412 "...and {pending_progress_bars} more not shown above.",
413 )
414 .expect("valid template"),
415 ),
416 TickSettings::default(),
417 );
418 let mp = pb_manager.mp.clone();
419
420 Self {
421 pb_manager: Mutex::new(pb_manager),
422 mp,
423 span_field_formatter: DefaultFields::new(),
424 progress_style: ProgressStyle::with_template(
425 "{span_child_prefix}{spinner} {span_name}{{{span_fields}}}",
426 )
427 .expect("valid template"),
428 span_child_prefix_indent: " ",
429 span_child_prefix_symbol: "↳ ",
430 get_context: WithContext(Self::get_context),
431 get_stderr_writer_context: WithStderrWriter(Self::get_stderr_writer_context),
432 get_stdout_writer_context: WithStdoutWriter(Self::get_stdout_writer_context),
433 get_multi_progress_context: WithMultiProgress(Self::get_multi_progress_context),
434 inner: PhantomData,
435 }
436 }
437}
438
439impl<S, F> IndicatifLayer<S, F>
441where
442 S: Subscriber + for<'a> LookupSpan<'a>,
443{
444 #[deprecated(since = "0.2.3", note = "use get_stderr_writer() instead")]
445 pub fn get_fmt_writer(&self) -> IndicatifWriter<writer::Stderr> {
446 self.get_stderr_writer()
447 }
448
449 pub fn get_stderr_writer(&self) -> IndicatifWriter<writer::Stderr> {
458 IndicatifWriter::new(self.mp.clone())
460 }
461
462 pub fn get_stdout_writer(&self) -> IndicatifWriter<writer::Stdout> {
471 IndicatifWriter::new(self.mp.clone())
473 }
474
475 pub fn with_span_field_formatter<F2>(self, formatter: F2) -> IndicatifLayer<S, F2>
480 where
481 F2: for<'writer> FormatFields<'writer> + 'static,
482 {
483 IndicatifLayer {
484 pb_manager: self.pb_manager,
485 mp: self.mp,
486 span_field_formatter: formatter,
487 progress_style: self.progress_style,
488 span_child_prefix_indent: self.span_child_prefix_indent,
489 span_child_prefix_symbol: self.span_child_prefix_symbol,
490 get_context: WithContext(IndicatifLayer::<S, F2>::get_context),
491 get_stderr_writer_context: WithStderrWriter(
492 IndicatifLayer::<S, F2>::get_stderr_writer_context,
493 ),
494 get_stdout_writer_context: WithStdoutWriter(
495 IndicatifLayer::<S, F2>::get_stdout_writer_context,
496 ),
497 get_multi_progress_context: WithMultiProgress(
498 IndicatifLayer::<S, F2>::get_multi_progress_context,
499 ),
500 inner: self.inner,
501 }
502 }
503
504 pub fn with_progress_style(mut self, style: ProgressStyle) -> Self {
514 self.progress_style = style;
515 self
516 }
517
518 pub fn with_span_child_prefix_indent(mut self, indent: &'static str) -> Self {
524 self.span_child_prefix_indent = indent;
525 self
526 }
527
528 pub fn with_span_child_prefix_symbol(mut self, symbol: &'static str) -> Self {
533 self.span_child_prefix_symbol = symbol;
534 self
535 }
536
537 pub fn with_max_progress_bars(
545 mut self,
546 max_progress_bars: u64,
547 footer_style: Option<ProgressStyle>,
548 ) -> Self {
549 if let Ok(pb_manager) = self.pb_manager.get_mut() {
550 pb_manager.set_max_progress_bars(max_progress_bars, footer_style);
551 }
552
553 self
554 }
555
556 pub fn with_tick_settings(mut self, tick_settings: TickSettings) -> Self {
558 if let Ok(pb_manager) = self.pb_manager.get_mut() {
559 pb_manager.set_tick_settings(tick_settings);
560 }
561
562 self
563 }
564}
565
566impl<S, F> IndicatifLayer<S, F>
567where
568 S: Subscriber + for<'a> LookupSpan<'a>,
569 F: for<'writer> FormatFields<'writer> + 'static,
570{
571 fn get_context(
572 dispatch: &tracing::Dispatch,
573 id: &span::Id,
574 f: &mut dyn FnMut(&mut IndicatifSpanContext),
575 ) {
576 let subscriber = dispatch
579 .downcast_ref::<S>()
580 .expect("subscriber should downcast to expected type; this is a bug!");
581 let span = subscriber
582 .span(id)
583 .expect("Span not found in context, this is a bug");
584
585 let mut ext = span.extensions_mut();
586
587 if let Some(indicatif_ctx) = ext.get_mut::<IndicatifSpanContext>() {
588 f(indicatif_ctx);
589 }
590 }
591
592 fn get_stderr_writer_context(
593 dispatch: &tracing::Dispatch,
594 f: &mut dyn FnMut(IndicatifWriter<writer::Stderr>),
595 ) {
596 let layer = dispatch
597 .downcast_ref::<IndicatifLayer<S, F>>()
598 .expect("subscriber should downcast to expected type; this is a bug!");
599
600 f(layer.get_stderr_writer())
601 }
602
603 fn get_stdout_writer_context(
604 dispatch: &tracing::Dispatch,
605 f: &mut dyn FnMut(IndicatifWriter<writer::Stdout>),
606 ) {
607 let layer = dispatch
608 .downcast_ref::<IndicatifLayer<S, F>>()
609 .expect("subscriber should downcast to expected type; this is a bug!");
610
611 f(layer.get_stdout_writer())
612 }
613
614 fn get_multi_progress_context(dispatch: &tracing::Dispatch, f: &mut dyn FnMut(MultiProgress)) {
615 let layer = dispatch
616 .downcast_ref::<IndicatifLayer<S, F>>()
617 .expect("subscriber should downcast to expected type; this is a bug!");
618
619 f(layer.mp.clone())
620 }
621
622 fn handle_on_enter(
623 &self,
624 pb_manager: &mut ProgressBarManager,
625 id: &span::Id,
626 ctx: &layer::Context<'_, S>,
627 ) -> Option<ProgressBar> {
628 let span = ctx
629 .span(id)
630 .expect("Span not found in context, this is a bug");
631 let mut ext = span.extensions_mut();
632
633 if let Some(indicatif_ctx) = ext.get_mut::<IndicatifSpanContext>() {
634 if indicatif_ctx.progress_bar.is_none() {
636 indicatif_ctx.make_progress_bar(&self.progress_style);
637
638 if let Some(ref parent_span_with_pb) = indicatif_ctx.parent_span {
639 let parent_pb = self.handle_on_enter(pb_manager, parent_span_with_pb, ctx);
641
642 indicatif_ctx.parent_progress_bar = parent_pb;
643 }
644
645 pb_manager.show_progress_bar(indicatif_ctx, id);
646 }
647
648 return indicatif_ctx.progress_bar.to_owned();
649 }
650
651 None
652 }
653}
654
655impl<S, F> layer::Layer<S> for IndicatifLayer<S, F>
656where
657 S: Subscriber + for<'a> LookupSpan<'a>,
658 F: for<'writer> FormatFields<'writer> + 'static,
659{
660 fn on_new_span(&self, attrs: &span::Attributes<'_>, id: &span::Id, ctx: layer::Context<'_, S>) {
661 let span = ctx
662 .span(id)
663 .expect("Span not found in context, this is a bug");
664 let mut ext = span.extensions_mut();
665
666 let mut fields = FormattedFields::<F>::new(String::new());
667 let _ = self
668 .span_field_formatter
669 .format_fields(fields.as_writer(), attrs);
670
671 let parent_span = ctx.span_scope(id).and_then(|scope| {
673 scope.skip(1).find(|span| {
674 let ext = span.extensions();
675
676 ext.get::<IndicatifSpanContext>().is_some()
677 })
678 });
679 let parent_span_id = parent_span.as_ref().map(|span| span.id());
680 let parent_span_ext = parent_span.as_ref().map(|span| span.extensions());
681 let parent_indicatif_ctx = parent_span_ext.as_ref().map(|ext| {
682 ext.get::<IndicatifSpanContext>()
683 .expect("validated it exists prior")
684 });
685
686 let (span_child_prefix, level) = match parent_indicatif_ctx {
687 Some(v) => {
688 let level = v.level + 1;
689
690 (
691 format!(
692 "{}{}",
693 self.span_child_prefix_indent.repeat(level.into()),
694 self.span_child_prefix_symbol
695 ),
696 level,
697 )
698 }
699 None => (String::new(), 0),
700 };
701
702 ext.insert(IndicatifSpanContext {
703 progress_bar: None,
704 pb_init_settings: ProgressBarInitSettings::default(),
705 parent_progress_bar: None,
706 parent_span: parent_span_id,
707 span_fields_formatted: Some(fields.fields),
708 span_name: span.name().to_string(),
709 span_child_prefix,
710 level,
711 finish_message: None,
712 });
713 }
714
715 fn on_enter(&self, id: &span::Id, ctx: layer::Context<'_, S>) {
716 if let Ok(mut pb_manager_lock) = self.pb_manager.lock() {
717 self.handle_on_enter(&mut pb_manager_lock, id, &ctx);
718 }
719 }
720
721 fn on_close(&self, id: span::Id, ctx: layer::Context<'_, S>) {
722 if let Ok(mut pb_manager_lock) = self.pb_manager.lock() {
723 let span = ctx
724 .span(&id)
725 .expect("Span not found in context, this is a bug");
726 let mut ext = span.extensions_mut();
727
728 if let Some(indicatif_ctx) = ext.get_mut::<IndicatifSpanContext>() {
730 pb_manager_lock.finish_progress_bar(indicatif_ctx, &ctx);
731 }
732 }
733 }
734
735 unsafe fn downcast_raw(&self, id: TypeId) -> Option<*const ()> {
740 match id {
741 id if id == TypeId::of::<Self>() => Some(self as *const _ as *const ()),
742 id if id == TypeId::of::<WithContext>() => {
743 Some(&self.get_context as *const _ as *const ())
744 }
745 id if id == TypeId::of::<WithStderrWriter>() => {
746 Some(&self.get_stderr_writer_context as *const _ as *const ())
747 }
748 id if id == TypeId::of::<WithStdoutWriter>() => {
749 Some(&self.get_stdout_writer_context as *const _ as *const ())
750 }
751 id if id == TypeId::of::<WithMultiProgress>() => {
752 Some(&self.get_multi_progress_context as *const _ as *const ())
753 }
754 _ => None,
755 }
756 }
757}
758
759pub fn suspend_tracing_indicatif<F: FnOnce() -> R, R>(f: F) -> R {
772 let mut mp: Option<MultiProgress> = None;
773
774 tracing::dispatcher::get_default(|dispatch| {
775 if let Some(ctx) = dispatch.downcast_ref::<WithMultiProgress>() {
776 ctx.with_context(dispatch, |fetched_mp| {
777 mp = Some(fetched_mp);
778 })
779 }
780 });
781
782 if let Some(mp) = mp {
783 mp.suspend(f)
784 } else {
785 f()
786 }
787}
788
789#[macro_export]
795macro_rules! indicatif_println {
796 ($($arg:tt)*) => {
797 {
798 use std::io::Write;
799
800 if let Some(mut writer) = $crate::writer::get_indicatif_stdout_writer() {
801 writeln!(writer, $($arg)*).unwrap();
802 } else {
803 #[allow(clippy::explicit_write)]
804 writeln!(std::io::stdout(), $($arg)*).unwrap();
805 }
806 }
807 }
808}
809
810#[macro_export]
816macro_rules! indicatif_eprintln {
817 ($($arg:tt)*) => {
818 {
819 use std::io::Write;
820
821 if let Some(mut writer) = $crate::writer::get_indicatif_stderr_writer() {
822 writeln!(writer, $($arg)*).unwrap();
823 } else {
824 #[allow(clippy::explicit_write)]
825 writeln!(std::io::stderr(), $($arg)*).unwrap();
826 }
827 }
828 }
829}
830
831#[cfg(test)]
832mod tests;