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

// twink, a little logging crate
// Copyright (c) 2023 fawn and rini
//
// SPDX-License-Identifier: Apache-2.0

use proc_macro::TokenStream;
use quote::quote;
use syn::{parse_macro_input, LitStr};

static COLOR_NAMES: &[&str] = &[
    "black", "red", "green", "yellow", "blue", "purple", "cyan", "white",
];

fn parse_color(c: &str) -> Option<String> {
    let n = match c {
        "/" => 0,
        "b" | "bold" => 1,
        "d" | "dim" => 2,
        "i" | "italic" => 3,
        "u" | "underline" => 4,
        "gray" => 90, // "bright black"
        mut tok => {
            let mut n = 30;
            if let Some(c) = tok.strip_prefix("on ") {
                n += 10;
                tok = c;
            }
            if let Some(c) = tok.strip_prefix("bright ") {
                n += 60;
                tok = c;
            }
            n + COLOR_NAMES.iter().position(|&x| x == tok)?
        }
    };

    Some(format!("\x1b[{n}m"))
}

#[proc_macro]
pub fn fmt(input: TokenStream) -> TokenStream {
    let mut fmt = String::new();

    for tok in input {
        let stream = TokenStream::from(tok);
        fmt.push_str(&parse_macro_input!(stream as LitStr).value());
    }

    let map = fmt
        .match_indices('<')
        .rev()
        .map(|(i, _)| i)
        .collect::<Vec<usize>>();
    for i in map {
        if let Some(j) = fmt[i..].find('>') {
            if let Some(color) = parse_color(&fmt[i + 1..=i + j - 1]) {
                fmt.replace_range(i..=i + j, &color);
            }
        }
    }

    TokenStream::from(quote! { #fmt })
}

#[cfg(test)]
mod tests {
    macro_rules! test {
        ($name:ident: $input:expr, $output:expr) => {
            #[test]
            fn $name() {
                assert_eq!(fmt!($input), $output);
            }
        };
    }

    test!(simple: "<red>fuming</>", "\x1b[31mfuming\x1b[0m");
    test!(on_bright: "<purple>meooooooow <b><red><on bright purple>:3</>", "\x1b[35mmeooooooow \x1b[1m\x1b[31m\x1b[105m:3\x1b[0m");
}