1use std::any::TypeId;
40use std::marker::PhantomData;
41use std::sync::Mutex;
42
43use indicatif::MultiProgress;
44use indicatif::ProgressBar;
45pub use indicatif::style;
47use indicatif::style::ProgressStyle;
48use indicatif::style::ProgressTracker;
49use tracing_core::Subscriber;
50use tracing_core::span;
51use tracing_subscriber::fmt::FormatFields;
52use tracing_subscriber::fmt::FormattedFields;
53use tracing_subscriber::fmt::format::DefaultFields;
54use tracing_subscriber::layer;
55use tracing_subscriber::registry::LookupSpan;
56
57pub mod filter;
58mod pb_manager;
59pub mod span_ext;
60pub mod util;
61pub mod writer;
62
63use pb_manager::ProgressBarManager;
64pub use pb_manager::TickSettings;
65#[doc(inline)]
66pub use writer::IndicatifWriter;
67
68#[derive(Clone)]
69struct IndicatifProgressKey {
70 message: String,
71}
72
73impl ProgressTracker for IndicatifProgressKey {
74 fn clone_box(&self) -> Box<dyn ProgressTracker> {
75 Box::new(self.clone())
76 }
77
78 fn tick(&mut self, _: &indicatif::ProgressState, _: std::time::Instant) {}
79
80 fn reset(&mut self, _: &indicatif::ProgressState, _: std::time::Instant) {}
81
82 fn write(&self, _: &indicatif::ProgressState, w: &mut dyn std::fmt::Write) {
83 let _ = w.write_str(&self.message);
84 }
85}
86
87#[allow(clippy::type_complexity)]
110pub(crate) struct WithContext(
111 fn(&tracing::Dispatch, &span::Id, f: &mut dyn FnMut(&mut IndicatifSpanContext)),
112);
113
114#[allow(clippy::type_complexity)]
115pub(crate) struct WithStderrWriter(
116 fn(&tracing::Dispatch, f: &mut dyn FnMut(IndicatifWriter<writer::Stderr>)),
117);
118
119#[allow(clippy::type_complexity)]
120pub(crate) struct WithStdoutWriter(
121 fn(&tracing::Dispatch, f: &mut dyn FnMut(IndicatifWriter<writer::Stdout>)),
122);
123
124#[allow(clippy::type_complexity)]
125pub(crate) struct WithMultiProgress(fn(&tracing::Dispatch, f: &mut dyn FnMut(MultiProgress)));
126
127impl WithContext {
128 pub(crate) fn with_context(
129 &self,
130 dispatch: &tracing::Dispatch,
131 id: &span::Id,
132 mut f: impl FnMut(&mut IndicatifSpanContext),
133 ) {
134 (self.0)(dispatch, id, &mut f)
135 }
136}
137
138impl WithStderrWriter {
139 pub(crate) fn with_context(
140 &self,
141 dispatch: &tracing::Dispatch,
142 mut f: impl FnMut(IndicatifWriter<writer::Stderr>),
143 ) {
144 (self.0)(dispatch, &mut f)
145 }
146}
147
148impl WithStdoutWriter {
149 pub(crate) fn with_context(
150 &self,
151 dispatch: &tracing::Dispatch,
152 mut f: impl FnMut(IndicatifWriter<writer::Stdout>),
153 ) {
154 (self.0)(dispatch, &mut f)
155 }
156}
157
158impl WithMultiProgress {
159 pub(crate) fn with_context(
160 &self,
161 dispatch: &tracing::Dispatch,
162 mut f: impl FnMut(MultiProgress),
163 ) {
164 (self.0)(dispatch, &mut f)
165 }
166}
167
168#[derive(Default)]
169struct ProgressBarInitSettings {
170 style: Option<ProgressStyle>,
171 len: Option<u64>,
172 pos: Option<u64>,
173 message: Option<String>,
174}
175
176struct IndicatifSpanContext {
177 progress_bar: Option<ProgressBar>,
182 pb_init_settings: ProgressBarInitSettings,
184 parent_progress_bar: Option<ProgressBar>,
189 parent_span: Option<span::Id>,
191 span_fields_formatted: Option<String>,
193 span_name: String,
194 span_child_prefix: String,
195 level: u16,
198 finish_message: Option<String>,
200}
201
202impl IndicatifSpanContext {
203 fn add_keys_to_style(&self, style: ProgressStyle) -> ProgressStyle {
204 style
205 .with_key(
206 "span_name",
207 IndicatifProgressKey {
208 message: self.span_name.clone(),
209 },
210 )
211 .with_key(
212 "span_fields",
213 IndicatifProgressKey {
214 message: self.span_fields_formatted.to_owned().unwrap_or_default(),
215 },
216 )
217 .with_key(
218 "span_child_prefix",
219 IndicatifProgressKey {
220 message: self.span_child_prefix.clone(),
221 },
222 )
223 }
224
225 fn make_progress_bar(&mut self, default_style: &ProgressStyle) {
226 if self.progress_bar.is_none() {
227 let pb = ProgressBar::hidden().with_style(
228 self.pb_init_settings
229 .style
230 .take()
231 .unwrap_or_else(|| self.add_keys_to_style(default_style.clone())),
232 );
233
234 if let Some(len) = self.pb_init_settings.len.take() {
235 pb.set_length(len);
236 }
237
238 if let Some(msg) = self.pb_init_settings.message.take() {
239 pb.set_message(msg);
240 }
241
242 if let Some(pos) = self.pb_init_settings.pos.take() {
243 pb.set_position(pos);
244 }
245
246 self.progress_bar = Some(pb);
247 }
248 }
249
250 fn set_progress_bar_style(&mut self, style: ProgressStyle) {
251 if let Some(ref pb) = self.progress_bar {
252 pb.set_style(self.add_keys_to_style(style));
253 } else {
254 self.pb_init_settings.style = Some(self.add_keys_to_style(style));
255 }
256 }
257
258 fn set_progress_bar_length(&mut self, len: u64) {
259 if let Some(ref pb) = self.progress_bar {
260 pb.set_length(len);
261 } else {
262 self.pb_init_settings.len = Some(len);
263 }
264 }
265
266 fn set_progress_bar_position(&mut self, pos: u64) {
267 if let Some(ref pb) = self.progress_bar {
268 pb.set_position(pos);
269 } else {
270 self.pb_init_settings.pos = Some(pos);
271 }
272 }
273
274 fn set_progress_bar_message(&mut self, msg: String) {
275 if let Some(ref pb) = self.progress_bar {
276 pb.set_message(msg);
277 } else {
278 self.pb_init_settings.message = Some(msg);
279 }
280 }
281
282 fn inc_progress_bar_position(&mut self, pos: u64) {
283 if let Some(ref pb) = self.progress_bar {
284 pb.inc(pos);
285 } else if let Some(ref mut pb_pos) = self.pb_init_settings.pos {
286 *pb_pos += pos;
287 } else {
288 self.pb_init_settings.pos = Some(pos);
290 }
291 }
292
293 fn inc_progress_bar_length(&mut self, len: u64) {
294 if let Some(ref pb) = self.progress_bar {
295 pb.inc_length(len);
296 } else if let Some(ref mut pb_len) = self.pb_init_settings.len {
297 *pb_len += len;
298 }
299 }
300
301 fn progress_bar_tick(&mut self) {
302 if let Some(ref pb) = self.progress_bar {
303 pb.tick()
304 }
305 }
306
307 fn reset_progress_bar(&mut self) {
308 if let Some(ref pb) = self.progress_bar {
309 pb.reset();
310 }
311 }
312
313 fn reset_progress_bar_elapsed(&mut self) {
314 if let Some(ref pb) = self.progress_bar {
315 pb.reset_elapsed();
316 }
317 }
318
319 fn reset_progress_bar_eta(&mut self) {
320 if let Some(ref pb) = self.progress_bar {
321 pb.reset_eta();
322 }
323 }
324
325 fn set_progress_bar_finish_message(&mut self, msg: String) {
326 self.finish_message = Some(msg);
327 }
328}
329
330pub struct IndicatifLayer<S, F = DefaultFields> {
348 pb_manager: Mutex<ProgressBarManager>,
349 mp: MultiProgress,
352 span_field_formatter: F,
353 progress_style: ProgressStyle,
354 span_child_prefix_indent: &'static str,
355 span_child_prefix_symbol: &'static str,
356 get_context: WithContext,
357 get_stderr_writer_context: WithStderrWriter,
358 get_stdout_writer_context: WithStdoutWriter,
359 get_multi_progress_context: WithMultiProgress,
360 inner: PhantomData<S>,
361}
362
363impl<S> IndicatifLayer<S>
364where
365 S: Subscriber + for<'a> LookupSpan<'a>,
366{
367 pub fn new() -> Self {
382 Self::default()
383 }
384}
385
386impl<S> Default for IndicatifLayer<S>
387where
388 S: Subscriber + for<'a> LookupSpan<'a>,
389{
390 fn default() -> Self {
391 let pb_manager = ProgressBarManager::new(
392 7,
393 Some(
394 ProgressStyle::with_template(
395 "...and {pending_progress_bars} more not shown above.",
396 )
397 .expect("valid template"),
398 ),
399 TickSettings::default(),
400 );
401 let mp = pb_manager.mp.clone();
402
403 Self {
404 pb_manager: Mutex::new(pb_manager),
405 mp,
406 span_field_formatter: DefaultFields::new(),
407 progress_style: ProgressStyle::with_template(
408 "{span_child_prefix}{spinner} {span_name}{{{span_fields}}}",
409 )
410 .expect("valid template"),
411 span_child_prefix_indent: " ",
412 span_child_prefix_symbol: "↳ ",
413 get_context: WithContext(Self::get_context),
414 get_stderr_writer_context: WithStderrWriter(Self::get_stderr_writer_context),
415 get_stdout_writer_context: WithStdoutWriter(Self::get_stdout_writer_context),
416 get_multi_progress_context: WithMultiProgress(Self::get_multi_progress_context),
417 inner: PhantomData,
418 }
419 }
420}
421
422impl<S, F> IndicatifLayer<S, F>
424where
425 S: Subscriber + for<'a> LookupSpan<'a>,
426{
427 #[deprecated(since = "0.2.3", note = "use get_stderr_writer() instead")]
428 pub fn get_fmt_writer(&self) -> IndicatifWriter<writer::Stderr> {
429 self.get_stderr_writer()
430 }
431
432 pub fn get_stderr_writer(&self) -> IndicatifWriter<writer::Stderr> {
441 IndicatifWriter::new(self.mp.clone())
443 }
444
445 pub fn get_stdout_writer(&self) -> IndicatifWriter<writer::Stdout> {
454 IndicatifWriter::new(self.mp.clone())
456 }
457
458 pub fn with_span_field_formatter<F2>(self, formatter: F2) -> IndicatifLayer<S, F2>
463 where
464 F2: for<'writer> FormatFields<'writer> + 'static,
465 {
466 IndicatifLayer {
467 pb_manager: self.pb_manager,
468 mp: self.mp,
469 span_field_formatter: formatter,
470 progress_style: self.progress_style,
471 span_child_prefix_indent: self.span_child_prefix_indent,
472 span_child_prefix_symbol: self.span_child_prefix_symbol,
473 get_context: WithContext(IndicatifLayer::<S, F2>::get_context),
474 get_stderr_writer_context: WithStderrWriter(
475 IndicatifLayer::<S, F2>::get_stderr_writer_context,
476 ),
477 get_stdout_writer_context: WithStdoutWriter(
478 IndicatifLayer::<S, F2>::get_stdout_writer_context,
479 ),
480 get_multi_progress_context: WithMultiProgress(
481 IndicatifLayer::<S, F2>::get_multi_progress_context,
482 ),
483 inner: self.inner,
484 }
485 }
486
487 pub fn with_progress_style(mut self, style: ProgressStyle) -> Self {
497 self.progress_style = style;
498 self
499 }
500
501 pub fn with_span_child_prefix_indent(mut self, indent: &'static str) -> Self {
507 self.span_child_prefix_indent = indent;
508 self
509 }
510
511 pub fn with_span_child_prefix_symbol(mut self, symbol: &'static str) -> Self {
516 self.span_child_prefix_symbol = symbol;
517 self
518 }
519
520 pub fn with_max_progress_bars(
528 mut self,
529 max_progress_bars: u64,
530 footer_style: Option<ProgressStyle>,
531 ) -> Self {
532 if let Ok(pb_manager) = self.pb_manager.get_mut() {
533 pb_manager.set_max_progress_bars(max_progress_bars, footer_style);
534 }
535
536 self
537 }
538
539 pub fn with_tick_settings(mut self, tick_settings: TickSettings) -> Self {
541 if let Ok(pb_manager) = self.pb_manager.get_mut() {
542 pb_manager.set_tick_settings(tick_settings);
543 }
544
545 self
546 }
547}
548
549impl<S, F> IndicatifLayer<S, F>
550where
551 S: Subscriber + for<'a> LookupSpan<'a>,
552 F: for<'writer> FormatFields<'writer> + 'static,
553{
554 fn get_context(
555 dispatch: &tracing::Dispatch,
556 id: &span::Id,
557 f: &mut dyn FnMut(&mut IndicatifSpanContext),
558 ) {
559 let subscriber = dispatch
562 .downcast_ref::<S>()
563 .expect("subscriber should downcast to expected type; this is a bug!");
564 let span = subscriber
565 .span(id)
566 .expect("Span not found in context, this is a bug");
567
568 let mut ext = span.extensions_mut();
569
570 if let Some(indicatif_ctx) = ext.get_mut::<IndicatifSpanContext>() {
571 f(indicatif_ctx);
572 }
573 }
574
575 fn get_stderr_writer_context(
576 dispatch: &tracing::Dispatch,
577 f: &mut dyn FnMut(IndicatifWriter<writer::Stderr>),
578 ) {
579 let layer = dispatch
580 .downcast_ref::<IndicatifLayer<S, F>>()
581 .expect("subscriber should downcast to expected type; this is a bug!");
582
583 f(layer.get_stderr_writer())
584 }
585
586 fn get_stdout_writer_context(
587 dispatch: &tracing::Dispatch,
588 f: &mut dyn FnMut(IndicatifWriter<writer::Stdout>),
589 ) {
590 let layer = dispatch
591 .downcast_ref::<IndicatifLayer<S, F>>()
592 .expect("subscriber should downcast to expected type; this is a bug!");
593
594 f(layer.get_stdout_writer())
595 }
596
597 fn get_multi_progress_context(dispatch: &tracing::Dispatch, f: &mut dyn FnMut(MultiProgress)) {
598 let layer = dispatch
599 .downcast_ref::<IndicatifLayer<S, F>>()
600 .expect("subscriber should downcast to expected type; this is a bug!");
601
602 f(layer.mp.clone())
603 }
604
605 fn handle_on_enter(
606 &self,
607 pb_manager: &mut ProgressBarManager,
608 id: &span::Id,
609 ctx: &layer::Context<'_, S>,
610 ) -> Option<ProgressBar> {
611 let span = ctx
612 .span(id)
613 .expect("Span not found in context, this is a bug");
614 let mut ext = span.extensions_mut();
615
616 if let Some(indicatif_ctx) = ext.get_mut::<IndicatifSpanContext>() {
617 if indicatif_ctx.progress_bar.is_none() {
619 indicatif_ctx.make_progress_bar(&self.progress_style);
620
621 if let Some(ref parent_span_with_pb) = indicatif_ctx.parent_span {
622 let parent_pb = self.handle_on_enter(pb_manager, parent_span_with_pb, ctx);
624
625 indicatif_ctx.parent_progress_bar = parent_pb;
626 }
627
628 pb_manager.show_progress_bar(indicatif_ctx, id);
629 }
630
631 return indicatif_ctx.progress_bar.to_owned();
632 }
633
634 None
635 }
636}
637
638impl<S, F> layer::Layer<S> for IndicatifLayer<S, F>
639where
640 S: Subscriber + for<'a> LookupSpan<'a>,
641 F: for<'writer> FormatFields<'writer> + 'static,
642{
643 fn on_new_span(&self, attrs: &span::Attributes<'_>, id: &span::Id, ctx: layer::Context<'_, S>) {
644 let span = ctx
645 .span(id)
646 .expect("Span not found in context, this is a bug");
647 let mut ext = span.extensions_mut();
648
649 let mut fields = FormattedFields::<F>::new(String::new());
650 let _ = self
651 .span_field_formatter
652 .format_fields(fields.as_writer(), attrs);
653
654 let parent_span = ctx.span_scope(id).and_then(|scope| {
656 scope.skip(1).find(|span| {
657 let ext = span.extensions();
658
659 ext.get::<IndicatifSpanContext>().is_some()
660 })
661 });
662 let parent_span_id = parent_span.as_ref().map(|span| span.id());
663 let parent_span_ext = parent_span.as_ref().map(|span| span.extensions());
664 let parent_indicatif_ctx = parent_span_ext.as_ref().map(|ext| {
665 ext.get::<IndicatifSpanContext>()
666 .expect("validated it exists prior")
667 });
668
669 let (span_child_prefix, level) = match parent_indicatif_ctx {
670 Some(v) => {
671 let level = v.level + 1;
672
673 (
674 format!(
675 "{}{}",
676 self.span_child_prefix_indent.repeat(level.into()),
677 self.span_child_prefix_symbol
678 ),
679 level,
680 )
681 }
682 None => (String::new(), 0),
683 };
684
685 ext.insert(IndicatifSpanContext {
686 progress_bar: None,
687 pb_init_settings: ProgressBarInitSettings::default(),
688 parent_progress_bar: None,
689 parent_span: parent_span_id,
690 span_fields_formatted: Some(fields.fields),
691 span_name: span.name().to_string(),
692 span_child_prefix,
693 level,
694 finish_message: None,
695 });
696 }
697
698 fn on_enter(&self, id: &span::Id, ctx: layer::Context<'_, S>) {
699 if let Ok(mut pb_manager_lock) = self.pb_manager.lock() {
700 self.handle_on_enter(&mut pb_manager_lock, id, &ctx);
701 }
702 }
703
704 fn on_close(&self, id: span::Id, ctx: layer::Context<'_, S>) {
705 if let Ok(mut pb_manager_lock) = self.pb_manager.lock() {
706 let span = ctx
707 .span(&id)
708 .expect("Span not found in context, this is a bug");
709 let mut ext = span.extensions_mut();
710
711 if let Some(indicatif_ctx) = ext.get_mut::<IndicatifSpanContext>() {
713 pb_manager_lock.finish_progress_bar(indicatif_ctx, &ctx);
714 }
715 }
716 }
717
718 unsafe fn downcast_raw(&self, id: TypeId) -> Option<*const ()> {
723 match id {
724 id if id == TypeId::of::<Self>() => Some(self as *const _ as *const ()),
725 id if id == TypeId::of::<WithContext>() => {
726 Some(&self.get_context as *const _ as *const ())
727 }
728 id if id == TypeId::of::<WithStderrWriter>() => {
729 Some(&self.get_stderr_writer_context as *const _ as *const ())
730 }
731 id if id == TypeId::of::<WithStdoutWriter>() => {
732 Some(&self.get_stdout_writer_context as *const _ as *const ())
733 }
734 id if id == TypeId::of::<WithMultiProgress>() => {
735 Some(&self.get_multi_progress_context as *const _ as *const ())
736 }
737 _ => None,
738 }
739 }
740}
741
742pub fn suspend_tracing_indicatif<F: FnOnce() -> R, R>(f: F) -> R {
755 let mut mp: Option<MultiProgress> = None;
756
757 tracing::dispatcher::get_default(|dispatch| {
758 if let Some(ctx) = dispatch.downcast_ref::<WithMultiProgress>() {
759 ctx.with_context(dispatch, |fetched_mp| {
760 mp = Some(fetched_mp);
761 })
762 }
763 });
764
765 if let Some(mp) = mp {
766 mp.suspend(f)
767 } else {
768 f()
769 }
770}
771
772#[macro_export]
778macro_rules! indicatif_println {
779 ($($arg:tt)*) => {
780 {
781 use std::io::Write;
782
783 if let Some(mut writer) = $crate::writer::get_indicatif_stdout_writer() {
784 writeln!(writer, $($arg)*).unwrap();
785 } else {
786 #[allow(clippy::explicit_write)]
787 writeln!(std::io::stdout(), $($arg)*).unwrap();
788 }
789 }
790 }
791}
792
793#[macro_export]
799macro_rules! indicatif_eprintln {
800 ($($arg:tt)*) => {
801 {
802 use std::io::Write;
803
804 if let Some(mut writer) = $crate::writer::get_indicatif_stderr_writer() {
805 writeln!(writer, $($arg)*).unwrap();
806 } else {
807 #[allow(clippy::explicit_write)]
808 writeln!(std::io::stderr(), $($arg)*).unwrap();
809 }
810 }
811 }
812}
813
814#[cfg(test)]
815mod tests;