colored_macro_impl/
lib.rs

1//! This crate implements the macro for `colored_macro` and should not be used directly.
2
3use ansi::RESET;
4use parser::Element;
5use proc_macro2::TokenStream;
6use syn::{
7    parse::{Parse, ParseStream, Result},
8    parse2,
9    token::Comma,
10    Expr, LitStr,
11};
12
13use crate::{ansi::style_to_ansi_code, parser::parse_tags};
14mod ansi;
15mod parser;
16#[cfg(test)]
17mod tests;
18
19#[doc(hidden)]
20pub fn colored_macro(item: TokenStream) -> Result<TokenStream> {
21    Ok(process(parse2(item)?))
22}
23
24/// A segment is like a token, it can either be optionally styled text or a style end tag.
25#[derive(Debug, PartialEq)]
26pub(crate) enum Segment {
27    /// (Style, Text)
28    Text(Option<String>, String),
29    /// (Style)
30    StyleEnd(String),
31}
32
33#[derive(Debug, PartialEq)]
34pub(crate) struct Colored {
35    /// The segments of the string literal.
36    /// A segment is like a token, it can either be optionally styled text or a style end tag.
37    pub segments: Vec<Segment>,
38    /// Arguments passed after the string literal (e.g. `colored!("Hello, {}!", "world")`, `"world"` is a format argument).
39    /// This only includes arguments that literally appear after the comma, not inline arguments like `{name}`.
40    pub format_args: Vec<String>,
41}
42
43impl Parse for Colored {
44    /// Parse the macro input into a `Colored` struct.
45    /// The input is in the form `"String <red>literal</red>", format_args*`.
46    fn parse(input: ParseStream) -> Result<Self> {
47        let mut segments = Vec::new();
48        let mut format_args = Vec::new();
49        let mut style = None;
50
51        for element in parse_tags(input.parse::<LitStr>()?.value()) {
52            match element {
53                Element::Start(tag_name) => {
54                    style = Some(tag_name.to_string());
55                }
56                Element::End(tag_name) => {
57                    segments.push(Segment::StyleEnd(tag_name.to_string()));
58                }
59                Element::Text(text) => {
60                    segments.push(Segment::Text(style.take(), text));
61                }
62            }
63        }
64
65        while !input.is_empty() {
66            input.parse::<Comma>()?;
67            let expr = input.parse::<Expr>()?;
68            format_args.push(quote::quote!(#expr).to_string());
69        }
70
71        Ok(Self {
72            segments,
73            format_args,
74        })
75    }
76}
77
78pub(crate) fn process(colored: Colored) -> TokenStream {
79    let mut fmt_string = String::new();
80    let mut no_color_string = String::new();
81    let mut style_stack = Vec::new();
82
83    for segment in colored.segments {
84        match segment {
85            Segment::Text(style, text) => {
86                if let Some(style) = style {
87                    let ansi_style = style_to_ansi_code(style);
88                    fmt_string.push_str(&format!("{}{}", ansi_style, text));
89                    no_color_string.push_str(&text);
90                    style_stack.push(ansi_style);
91                } else {
92                    fmt_string.push_str(&text);
93                    no_color_string.push_str(&text);
94                }
95            }
96            Segment::StyleEnd(style) => {
97                let ansi_style = style_to_ansi_code(style);
98
99                if let Some(prev_ansi_style) = style_stack.pop() {
100                    if prev_ansi_style != ansi_style {
101                        panic!(
102                            "Mismatched style end tag: expected {}, got {}",
103                            prev_ansi_style.escape_default(),
104                            ansi_style.escape_default()
105                        );
106                    }
107                }
108
109                let reset = if let Some(prev_ansi_style) = style_stack.last() {
110                    format!("\x1b[0;00m{}", prev_ansi_style)
111                } else {
112                    RESET.to_string()
113                };
114
115                fmt_string.push_str(&reset);
116            }
117        }
118    }
119
120    fmt_string.push_str(RESET);
121
122    let mut output = String::new();
123
124    if cfg!(feature = "no-color") {
125        output.push_str("if std::env::var(\"NO_COLOR\").map(|v| v == \"1\").unwrap_or(false) { ");
126        // Open a `format!("` call.
127        output.push_str("format!(\"");
128        output.push_str(&no_color_string);
129        // Close the `format!("` call with `"`.
130        output.push_str("\", ");
131
132        // Push the format arguments.
133        for arg in &colored.format_args {
134            output.push_str(arg);
135            output.push_str(", ");
136        }
137
138        output.push_str(") } else {");
139
140        // Open a `format!("` call.
141        output.push_str("format!(\"");
142        output.push_str(&fmt_string);
143        // Close the `format!("` call with `"`.
144        output.push_str("\", ");
145
146        // Push the format arguments.
147        for arg in &colored.format_args {
148            output.push_str(arg);
149            output.push_str(", ");
150        }
151
152        // Close the `format!` call
153        output.push_str(") }");
154    } else {
155        // Open a `format!("` call.
156        output.push_str("format!(\"");
157        output.push_str(&fmt_string);
158        // Close the `format!("` call with `"`.
159        output.push_str("\", ");
160
161        // Push the format arguments.
162        for arg in colored.format_args {
163            output.push_str(&arg);
164            output.push_str(", ");
165        }
166
167        // Close the `format!` call
168        output.push(')');
169    }
170
171    output.parse().unwrap()
172}