1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97

//! Provides a rudimentary filter layer that can be used to selectively enable progress bars on a
//! per-span level.
//!
//! # Example Use
//!
//! ```
//! use tracing_subscriber::layer::SubscriberExt;
//! use tracing_subscriber::util::SubscriberInitExt;
//! use tracing_indicatif::IndicatifLayer;
//! use tracing_indicatif::filter::IndicatifFilter;
//! use tracing_indicatif::filter::hide_indicatif_span_fields;
//! use tracing_subscriber::fmt::format::DefaultFields;
//! use tracing_subscriber::layer::Layer;
//!
//! let indicatif_layer = IndicatifLayer::new()
//!     .with_span_field_formatter(hide_indicatif_span_fields(DefaultFields::new()));
//!
//! tracing_subscriber::registry()
//!     .with(tracing_subscriber::fmt::layer().with_writer(indicatif_layer.get_stderr_writer()))
//!     .with(indicatif_layer.with_filter(IndicatifFilter::new(false)))
//!     .init();
//! ```
use std::fmt;
use std::marker::PhantomData;

use tracing_core::{Field, Subscriber};
use tracing_subscriber::layer::Filter;
use tracing_subscriber::{
    field::{MakeVisitor, VisitFmt, VisitOutput},
    fmt::format::Writer,
};

use crate::util::FilteredFormatFields;

/// A filter that filters based on the presence of a field with the name of either
/// "indicatif.pb_show" or "indicatif.pb_hide" on the span.
///
/// The value for this field is irrelevant and not factored in to the filtering (this is due to
/// tracing not making field values available in the `on_new_span` method). To avoid confusion, it
/// is recommended to set the value of this field to [tracing::field::Empty].
///
/// If both "indicatif.pb_show" and "indicatif.pb_hide" are present, the behavior is to show a
/// progress bar.
pub struct IndicatifFilter<S> {
    show_progress_bars_by_default: bool,
    subscriber: PhantomData<S>,
}

impl<S: Subscriber> IndicatifFilter<S> {
    /// Constructs the filter.
    ///
    /// If "indicatif.pb_show" or "indicatif.pb_hide" are not present as a field on the span,
    /// then the value of `show_progress_bars_by_default` is used; i.e. if
    /// `show_progress_bars_by_default` is `false`, then progress bars are not shown for spans by
    /// default.
    pub fn new(show_progress_bars_by_default: bool) -> Self {
        Self {
            show_progress_bars_by_default,
            subscriber: PhantomData,
        }
    }
}

impl<S: Subscriber> Filter<S> for IndicatifFilter<S> {
    fn enabled(
        &self,
        meta: &tracing::Metadata<'_>,
        _: &tracing_subscriber::layer::Context<'_, S>,
    ) -> bool {
        if !meta.is_span() {
            return false;
        }

        if meta.fields().field("indicatif.pb_show").is_some() {
            return true;
        }

        if meta.fields().field("indicatif.pb_hide").is_some() {
            return false;
        }

        self.show_progress_bars_by_default
    }
}

/// Returns a [tracing_subscriber::fmt::FormatFields] that ignores the "indicatif.pb_show" and "indicatif.pb_hide" fields.
pub fn hide_indicatif_span_fields<'writer, Format>(
    format: Format,
) -> FilteredFormatFields<Format, impl Fn(&Field) -> bool + Clone>
where
    Format: MakeVisitor<Writer<'writer>>,
    Format::Visitor: VisitFmt + VisitOutput<fmt::Result>,
{
    FilteredFormatFields::new(format, |field: &Field| {
        field.name() != "indicatif.pb_show" && field.name() != "indicatif.pb_hide"
    })
}