Skip to main content

fancy_tree/color/
mod.rs

1//! This module provides utilities for colorization.
2pub use choice::ColorChoice;
3use either::{Either, Left, Right};
4use mlua::{FromLua, IntoLua, Lua};
5use owo_colors::{
6    AnsiColors::{
7        self, Black, Blue, BrightBlack, BrightBlue, BrightCyan, BrightGreen, BrightMagenta,
8        BrightRed, BrightWhite, BrightYellow, Cyan, Green, Magenta, Red, White, Yellow,
9    },
10    DynColors,
11};
12
13mod choice;
14
15/// Either ANSI colors or full RGB.
16#[derive(Debug, Clone, Copy, PartialEq, Eq)]
17pub enum Color {
18    Ansi(AnsiColors),
19    Rgb(u8, u8, u8),
20}
21
22impl Color {
23    /// Maps ansi color names to their values.
24    const ANSI_NAME_MAP: [(&'static str, AnsiColors); 16] = [
25        ("black", Black),
26        ("red", Red),
27        ("green", Green),
28        ("yellow", Yellow),
29        ("blue", Blue),
30        ("magenta", Magenta),
31        ("cyan", Cyan),
32        ("white", White),
33        ("bright-black", BrightBlack),
34        ("bright-red", BrightRed),
35        ("bright-green", BrightGreen),
36        ("bright-yellow", BrightYellow),
37        ("bright-blue", BrightBlue),
38        ("bright-magenta", BrightMagenta),
39        ("bright-cyan", BrightCyan),
40        ("bright-white", BrightWhite),
41    ];
42
43    /// Tries to create an Ansi color from Lua.
44    fn ansi_from_lua_string(type_name: &'static str, s: &str) -> mlua::Result<Self> {
45        Self::ANSI_NAME_MAP
46            .into_iter()
47            .find_map(|(key, value)| (key == s).then_some(value))
48            .map(Self::Ansi)
49            .ok_or(mlua::Error::FromLuaConversionError {
50                from: type_name,
51                to: String::from("Color"),
52                message: Some(String::from("Expected one of the ansi color names")),
53            })
54    }
55
56    /// Tries to create an Rgb color from Lua.
57    fn rgb_from_lua_table(t: mlua::Table) -> mlua::Result<Self> {
58        // let [r, g, b] = ["r", "g", "b"].map(|key| t.get::<u8>(key));
59        t.get::<u8>("r")
60            .and_then(|r| t.get::<u8>("g").map(|g| (r, g)))
61            .and_then(|(r, g)| t.get::<u8>("b").map(|b| (r, g, b)))
62            .map(|(r, g, b)| Self::Rgb(r, g, b))
63    }
64
65    /// Gets the key name for the ANSI color.
66    fn ansi_name(ansi_colors: AnsiColors) -> &'static str {
67        debug_assert!(
68            Self::ANSI_NAME_MAP
69                .into_iter()
70                .any(|(_, color)| color == ansi_colors)
71        );
72
73        Self::ANSI_NAME_MAP
74            .into_iter()
75            .find_map(|(key, value)| (value == ansi_colors).then_some(key))
76            .expect("The mapping should exist")
77    }
78
79    /// Converts RGB into a table.
80    #[inline]
81    fn rgb_to_table(lua: &Lua, r: u8, g: u8, b: u8) -> mlua::Result<mlua::Table> {
82        lua.create_table_from([("r", r), ("g", g), ("b", b)])
83    }
84}
85
86impl FromLua for Color {
87    fn from_lua(value: mlua::Value, lua: &Lua) -> mlua::Result<Self> {
88        type AnsiOrRgb = Either<String, mlua::Table>;
89        let type_name = value.type_name();
90        let ansi_or_rgb = AnsiOrRgb::from_lua(value, lua)?;
91
92        match ansi_or_rgb {
93            Left(ansi) => Self::ansi_from_lua_string(type_name, &ansi),
94            Right(table) => Self::rgb_from_lua_table(table),
95        }
96    }
97}
98
99impl IntoLua for Color {
100    fn into_lua(self, lua: &Lua) -> mlua::Result<mlua::Value> {
101        match self {
102            Color::Ansi(ansi_colors) => Color::ansi_name(ansi_colors).into_lua(lua),
103            Color::Rgb(r, g, b) => Color::rgb_to_table(lua, r, g, b)?.into_lua(lua),
104        }
105    }
106}
107
108impl From<AnsiColors> for Color {
109    #[inline]
110    fn from(value: AnsiColors) -> Self {
111        Self::Ansi(value)
112    }
113}
114
115impl From<Color> for DynColors {
116    fn from(value: Color) -> Self {
117        match value {
118            Color::Ansi(color) => Self::Ansi(color),
119            Color::Rgb(r, g, b) => Self::Rgb(r, g, b),
120        }
121    }
122}
123
124#[cfg(test)]
125mod tests {
126    use super::*;
127    use rstest::rstest;
128
129    #[rstest]
130    #[case("black", AnsiColors::Black)]
131    #[case("magenta", AnsiColors::Magenta)]
132    #[case("bright-red", AnsiColors::BrightRed)]
133    #[case("bright-yellow", AnsiColors::BrightYellow)]
134    fn test_from_lua_string_ok(#[case] raw: &str, #[case] expected_ansi: AnsiColors) {
135        let lua = Lua::new();
136        let value = lua.create_string(raw).expect("A string to be created");
137        let value = mlua::Value::String(value);
138        let color = Color::from_lua(value, &lua).expect("Color should be converted");
139        let Color::Ansi(ansi) = color else {
140            panic!("Expected Color::Ansi")
141        };
142        assert_eq!(expected_ansi, ansi);
143    }
144
145    #[test]
146    fn test_from_lua_string_err() {
147        let lua = Lua::new();
148        let value = lua
149            .create_string("??unused??")
150            .expect("A string to be created");
151        let value = mlua::Value::String(value);
152        assert!(Color::from_lua(value, &lua).is_err());
153    }
154
155    #[test]
156    fn test_from_lua_tuple_ok() {
157        let lua = Lua::new();
158        let value = lua
159            .create_table_from([("r", 255u8), ("g", 0), ("b", 128)])
160            .expect("A table should be created");
161        let value = mlua::Value::Table(value);
162        let color = Color::from_lua(value, &lua).expect("Color should be converted");
163        let Color::Rgb(r, g, b) = color else {
164            panic!("Expected Color::Rgb")
165        };
166        assert_eq!(255, r);
167        assert_eq!(0, g);
168        assert_eq!(128, b);
169    }
170
171    #[test]
172    fn test_from_lua_tuple_err() {
173        let lua = Lua::new();
174        let value = lua
175            .create_table_from([("r", 255u8), ("b", 0)])
176            .expect("A table should be created");
177        let value = mlua::Value::Table(value);
178        assert!(Color::from_lua(value, &lua).is_err());
179    }
180}