inoculate/
trace.rs

1use crate::{
2    term::{style, ColorMode, OutputOptions, OwoColorize, Style},
3    Result,
4};
5use heck::TitleCase;
6use std::fmt;
7use tracing::{field::Field, Event, Level, Subscriber};
8use tracing_subscriber::{
9    field::Visit,
10    fmt::{format::Writer, FmtContext, FormatEvent, FormatFields, FormattedFields},
11    registry::LookupSpan,
12};
13
14impl OutputOptions {
15    pub fn trace_init(&self) -> Result<()> {
16        use tracing_subscriber::prelude::*;
17        let fmt = tracing_subscriber::fmt::layer()
18            .event_format(CargoFormatter {
19                styles: Styles::new(self.color),
20            })
21            .with_writer(std::io::stderr);
22
23        tracing_subscriber::registry()
24            .with(fmt)
25            .with(tracing_error::ErrorLayer::default())
26            .with(self.log.clone())
27            .try_init()?;
28        Ok(())
29    }
30}
31
32#[derive(Debug)]
33struct CargoFormatter {
34    styles: Styles,
35}
36
37struct Visitor<'styles, 'writer> {
38    level: Level,
39    writer: Writer<'writer>,
40    is_empty: bool,
41    styles: &'styles Styles,
42    did_cargo_format: bool,
43}
44
45#[derive(Debug)]
46struct Styles {
47    error: Style,
48    warn: Style,
49    info: Style,
50    debug: Style,
51    trace: Style,
52    pipes: Style,
53    bold: Style,
54}
55
56struct Prefixed<T> {
57    prefix: &'static str,
58    val: T,
59}
60
61impl<S, N> FormatEvent<S, N> for CargoFormatter
62where
63    S: Subscriber + for<'a> LookupSpan<'a>,
64    N: for<'a> FormatFields<'a> + 'static,
65{
66    fn format_event(
67        &self,
68        ctx: &FmtContext<'_, S, N>,
69        mut writer: Writer<'_>,
70        event: &Event<'_>,
71    ) -> fmt::Result {
72        let metadata = event.metadata();
73        let level = metadata.level();
74
75        let include_spans = {
76            let mut visitor = self.visitor(*level, writer.by_ref());
77            event.record(&mut visitor);
78            !visitor.did_cargo_format && ctx.lookup_current().is_some()
79        };
80
81        writer.write_char('\n')?;
82
83        if include_spans {
84            writeln!(
85                writer,
86                "   {} {}{}",
87                "-->".style(self.styles.pipes),
88                metadata.file().unwrap_or_else(|| metadata.target()),
89                DisplayOpt(metadata.line().map(Prefixed::prefix(":"))),
90            )?;
91            ctx.visit_spans(|span| {
92                let exts = span.extensions();
93                let fields = exts
94                    .get::<FormattedFields<N>>()
95                    .map(|f| f.fields.as_str())
96                    .unwrap_or("");
97                writeln!(
98                    writer,
99                    "    {} {}{}{}",
100                    "|".style(self.styles.pipes),
101                    span.name().style(self.styles.bold),
102                    if fields.is_empty() { "" } else { ": " },
103                    fields
104                )
105            })?;
106
107            writer.write_char('\n')?;
108        }
109
110        Ok(())
111    }
112}
113
114impl CargoFormatter {
115    fn visitor<'styles, 'writer>(
116        &'styles self,
117        level: Level,
118        writer: Writer<'writer>,
119    ) -> Visitor<'styles, 'writer> {
120        Visitor {
121            level,
122            writer,
123            is_empty: true,
124            styles: &self.styles,
125            did_cargo_format: false,
126        }
127    }
128}
129
130// === impl Visitor ===
131
132impl Visitor<'_, '_> {
133    const MESSAGE: &'static str = "message";
134    const INDENT: usize = 12;
135}
136
137impl Visit for Visitor<'_, '_> {
138    fn record_debug(&mut self, field: &Field, value: &dyn fmt::Debug) {
139        // If we're writing the first field of the event, either emit cargo
140        // formatting, or a level header.
141        if self.is_empty {
142            // If the level is `INFO` and it has a message that's
143            // shaped like a cargo log tag, emit the cargo tag followed by the
144            // rest of the message.
145            if self.level == Level::INFO && field.name() == Self::MESSAGE {
146                let message = format!("{value:?}");
147                if let Some((tag, message)) = message.as_str().split_once(' ') {
148                    if tag.len() <= Self::INDENT {
149                        let tag = tag.to_title_case();
150                        let style = match self.level {
151                            Level::DEBUG => self.styles.debug,
152                            _ => self.styles.info,
153                        };
154
155                        let _ = write!(
156                            self.writer,
157                            "{:>indent$} ",
158                            tag.style(style),
159                            indent = Self::INDENT
160                        );
161
162                        let _ = self.writer.write_str(message);
163                        self.is_empty = false;
164                        self.did_cargo_format = true;
165                        return;
166                    }
167                }
168            }
169
170            // Otherwise, emit a level tag.
171            let _ = match self.level {
172                Level::ERROR => write!(
173                    self.writer,
174                    "{}{} ",
175                    "error".style(self.styles.error),
176                    ":".style(self.styles.bold)
177                ),
178                Level::WARN => write!(
179                    self.writer,
180                    "{}{} ",
181                    "warning".style(self.styles.warn),
182                    ":".style(self.styles.bold),
183                ),
184                Level::INFO => write!(
185                    self.writer,
186                    "{}{} ",
187                    "info".style(self.styles.info),
188                    ":".style(self.styles.bold)
189                ),
190                Level::DEBUG => write!(
191                    self.writer,
192                    "{}{} ",
193                    "debug".style(self.styles.debug),
194                    ":".style(self.styles.bold)
195                ),
196                Level::TRACE => write!(
197                    self.writer,
198                    "{}{} ",
199                    "trace".style(self.styles.trace),
200                    ":".style(self.styles.bold)
201                ),
202            };
203        } else {
204            // If this is *not* the first field of the event, prefix it with a
205            // comma for the preceding field, instead of a cargo tag or level tag.
206            let _ = self.writer.write_str(", ");
207        }
208
209        if field.name() == Self::MESSAGE {
210            let _ = write!(self.writer, "{:?}", value.style(self.styles.bold));
211        } else {
212            let _ = write!(
213                self.writer,
214                "{}{} {:?}",
215                field.name().style(self.styles.bold),
216                ":".style(self.styles.bold),
217                value
218            );
219        }
220
221        self.is_empty = false;
222    }
223}
224
225// === impl Styles ===
226
227impl Styles {
228    fn new(colors: ColorMode) -> Self {
229        Self {
230            error: colors.if_color(style().red().bold()),
231            warn: colors.if_color(style().yellow().bold()),
232            info: colors.if_color(style().green().bold()),
233            debug: colors.if_color(style().blue().bold()),
234            trace: colors.if_color(style().purple().bold()),
235            bold: colors.if_color(style().bold()),
236            pipes: colors.if_color(style().blue().bold()),
237        }
238    }
239}
240
241impl<T> Prefixed<T> {
242    fn prefix(prefix: &'static str) -> impl Fn(T) -> Prefixed<T> {
243        move |val| Prefixed { val, prefix }
244    }
245}
246
247impl<T> fmt::Display for Prefixed<T>
248where
249    T: fmt::Display,
250{
251    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
252        write!(f, "{}{}", self.prefix, self.val)
253    }
254}
255
256impl<T> fmt::Debug for Prefixed<T>
257where
258    T: fmt::Debug,
259{
260    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
261        write!(f, "{}{:?}", self.prefix, self.val)
262    }
263}
264
265struct DisplayOpt<T>(Option<T>);
266
267impl<T> fmt::Display for DisplayOpt<T>
268where
269    T: fmt::Display,
270{
271    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
272        if let Some(ref val) = self.0 {
273            fmt::Display::fmt(val, f)?;
274        }
275
276        Ok(())
277    }
278}