feat: add user-defined color palette (#4209)

* docs(config): add color palette to docs

* feat: add user-defined color palette

* fix: update config schema

* refactor: apply suggestions from code review

Co-authored-by: David Knaack <davidkna@users.noreply.github.com>

* fix: update new test

* feat: add support for multiple palettes

* docs(config): update docs for multiple color palettes

* docs(config): fix formatting

* test: test overriding a predefined color with itself

* docs: mention palettes cannot reference themselves

* refactor: warn when using a nonexistent palette

* test: test retrieving a nonexistent color palette

* fix: fix issues with palette log messages

* fix: update config schema

* fix: skip serializing palette if none

* refactor: change nonexistent palette message to warning

* fix: update config schema

Co-authored-by: David Knaack <davidkna@users.noreply.github.com>

6 files changed, 169 insertions, 15 deletions

diff --git a/.github/config-schema.json b/.github/config-schema.json
index b04c593fc..68cdb9572 100644
--- a/.github/config-schema.json
+++ b/.github/config-schema.json
@@ -1544,6 +1544,22 @@
     "add_newline": {
       "default": true,
       "type": "boolean"
+    },
+    "palette": {
+      "type": [
+        "string",
+        "null"
+      ]
+    },
+    "palettes": {
+      "default": {},
+      "type": "object",
+      "additionalProperties": {
+        "type": "object",
+        "additionalProperties": {
+          "type": "string"
+        }
+      }
     }
   },
   "additionalProperties": false,
diff --git a/docs/config/README.md b/docs/config/README.md
index 93a60359a..85ddd0d18 100644
--- a/docs/config/README.md
+++ b/docs/config/README.md
@@ -176,13 +176,15 @@ This is the list of prompt-wide configuration options.
 
 ### Options
 
-| Option            | Default                        | Description                                                      |
-| ----------------- | ------------------------------ | ---------------------------------------------------------------- |
-| `format`          | [link](#default-prompt-format) | Configure the format of the prompt.                              |
-| `right_format`    | `""`                           | See [Enable Right Prompt](/advanced-config/#enable-right-prompt) |
-| `scan_timeout`    | `30`                           | Timeout for starship to scan files (in milliseconds).            |
-| `command_timeout` | `500`                          | Timeout for commands executed by starship (in milliseconds).     |
-| `add_newline`     | `true`                         | Inserts blank line between shell prompts.                        |
+| Option            | Default                        | Description                                                                                                                                                                      |
+| ----------------- | ------------------------------ | -------------------------------------------------------------------------------------------------------------------------------------------------------------------------------- |
+| `format`          | [link](#default-prompt-format) | Configure the format of the prompt.                                                                                                                                              |
+| `right_format`    | `""`                           | See [Enable Right Prompt](/advanced-config/#enable-right-prompt)                                                                                                                 |
+| `scan_timeout`    | `30`                           | Timeout for starship to scan files (in milliseconds).                                                                                                                            |
+| `command_timeout` | `500`                          | Timeout for commands executed by starship (in milliseconds).                                                                                                                     |
+| `add_newline`     | `true`                         | Inserts blank line between shell prompts.                                                                                                                                        |
+| `palette`         | `""`                           | Sets which color palette from `palettes` to use.                                                                                                                                 |
+| `palettes`        | `{}`                           | Collection of color palettes that assign [colors](/advanced-config/#style-strings) to user-defined names. Note that color palettes cannot reference their own color definitions. |
 
 ### Example
 
@@ -200,6 +202,16 @@ scan_timeout = 10
 
 # Disable the blank line at the start of the prompt
 add_newline = false
+
+# Set "foo" as custom color palette
+palette = "foo"
+
+# Define custom colors
+[palettes.foo]
+# Overwrite existing color
+blue = "21"
+# Define new color
+mustard = "#af8700"
 ```
 
 ### Default Prompt Format
diff --git a/src/config.rs b/src/config.rs
index 5f4ead789..7a44a3115 100644
--- a/src/config.rs
+++ b/src/config.rs
@@ -1,3 +1,5 @@
+use crate::configs::Palette;
+use crate::context::Context;
 use crate::serde_utils::ValueDeserializer;
 use crate::utils;
 use nu_ansi_term::Color;
@@ -7,6 +9,7 @@ use serde::{
 
 use std::borrow::Cow;
 use std::clone::Clone;
+use std::collections::HashMap;
 use std::io::ErrorKind;
 
 use std::env;
@@ -260,7 +263,7 @@ where
     D: Deserializer<'de>,
 {
     Cow::<'_, str>::deserialize(de).and_then(|s| {
-        parse_style_string(s.as_ref()).ok_or_else(|| D::Error::custom("Invalid style string"))
+        parse_style_string(s.as_ref(), None).ok_or_else(|| D::Error::custom("Invalid style string"))
     })
 }
 
@@ -275,7 +278,10 @@ where
  - 'blink'
  - '<color>'       (see the `parse_color_string` doc for valid color strings)
 */
-pub fn parse_style_string(style_string: &str) -> Option<nu_ansi_term::Style> {
+pub fn parse_style_string(
+    style_string: &str,
+    context: Option<&Context>,
+) -> Option<nu_ansi_term::Style> {
     style_string
         .split_whitespace()
         .fold(Some(nu_ansi_term::Style::new()), |maybe_style, token| {
@@ -308,7 +314,15 @@ pub fn parse_style_string(style_string: &str) -> Option<nu_ansi_term::Style> {
                             None // fg:none yields no style.
                         } else {
                             // Either bg or valid color or both.
-                            let parsed = parse_color_string(color_string);
+                            let parsed = parse_color_string(
+                                color_string,
+                                context.and_then(|x| {
+                                    get_palette(
+                                        &x.root_config.palettes,
+                                        x.root_config.palette.as_deref(),
+                                    )
+                                }),
+                            );
                             // bg + invalid color = reset the background to default.
                             if !col_fg && parsed.is_none() {
                                 let mut new_style = style;
@@ -335,9 +349,12 @@ pub fn parse_style_string(style_string: &str) -> Option<nu_ansi_term::Style> {
  There are three valid color formats:
   - #RRGGBB      (a hash followed by an RGB hex)
   - u8           (a number from 0-255, representing an ANSI color)
-  - colstring    (one of the 16 predefined color strings)
+  - colstring    (one of the 16 predefined color strings or a custom user-defined color)
 */
-fn parse_color_string(color_string: &str) -> Option<nu_ansi_term::Color> {
+fn parse_color_string(
+    color_string: &str,
+    palette: Option<&Palette>,
+) -> Option<nu_ansi_term::Color> {
     // Parse RGB hex values
     log::trace!("Parsing color_string: {}", color_string);
     if color_string.starts_with('#') {
@@ -362,6 +379,16 @@ fn parse_color_string(color_string: &str) -> Option<nu_ansi_term::Color> {
         return Some(Color::Fixed(ansi_color_num));
     }
 
+    // Check palette for a matching user-defined color
+    if let Some(palette_color) = palette.as_ref().and_then(|x| x.get(color_string)) {
+        log::trace!(
+            "Read user-defined color string: {} defined as {}",
+            color_string,
+            palette_color
+        );
+        return parse_color_string(palette_color, None);
+    }
+
     // Check for any predefined color strings
     // There are no predefined enums for bright colors, so we use Color::Fixed
     let predefined_color = match color_string.to_lowercase().as_str() {
@@ -392,6 +419,24 @@ fn parse_color_string(color_string: &str) -> Option<nu_ansi_term::Color> {
     predefined_color
 }
 
+fn get_palette<'a>(
+    palettes: &'a HashMap<String, Palette>,
+    palette_name: Option<&str>,
+) -> Option<&'a Palette> {
+    if let Some(palette_name) = palette_name {
+        let palette = palettes.get(palette_name);
+        if palette.is_some() {
+            log::trace!("Found color palette: {}", palette_name);
+        } else {
+            log::warn!("Could not find color palette: {}", palette_name);
+        }
+        palette
+    } else {
+        log::trace!("No color palette specified, using defaults");
+        None
+    }
+}
+
 #[cfg(test)]
 mod tests {
     use super::*;
@@ -778,4 +823,76 @@ mod tests {
             Style::new().fg(Color::Fixed(125)).on(Color::Fixed(127))
         );
     }
+
+    #[test]
+    fn table_get_colors_palette() {
+        // Test using colors defined in palette
+        let mut palette = Palette::new();
+        palette.insert("mustard".to_string(), "#af8700".to_string());
+        palette.insert("sky-blue".to_string(), "51".to_string());
+        palette.insert("red".to_string(), "#d70000".to_string());
+        palette.insert("blue".to_string(), "17".to_string());
+        palette.insert("green".to_string(), "green".to_string());
+
+        assert_eq!(
+            parse_color_string("mustard", Some(&palette)),
+            Some(Color::Rgb(175, 135, 0))
+        );
+        assert_eq!(
+            parse_color_string("sky-blue", Some(&palette)),
+            Some(Color::Fixed(51))
+        );
+
+        // Test overriding predefined colors
+        assert_eq!(
+            parse_color_string("red", Some(&palette)),
+            Some(Color::Rgb(215, 0, 0))
+        );
+        assert_eq!(
+            parse_color_string("blue", Some(&palette)),
+            Some(Color::Fixed(17))
+        );
+
+        // Test overriding a predefined color with itself
+        assert_eq!(
+            parse_color_string("green", Some(&palette)),
+            Some(Color::Green)
+        )
+    }
+
+    #[test]
+    fn table_get_palette() {
+        // Test retrieving color palette by name
+        let mut palette1 = Palette::new();
+        palette1.insert("test-color".to_string(), "123".to_string());
+
+        let mut palette2 = Palette::new();
+        palette2.insert("test-color".to_string(), "#ABCDEF".to_string());
+
+        let mut palettes = HashMap::<String, Palette>::new();
+        palettes.insert("palette1".to_string(), palette1);
+        palettes.insert("palette2".to_string(), palette2);
+
+        assert_eq!(
+            get_palette(&palettes, Some("palette1"))
+                .unwrap()
+                .get("test-color")
+                .unwrap(),
+            "123"
+        );
+
+        assert_eq!(
+            get_palette(&palettes, Some("palette2"))
+                .unwrap()
+                .get("test-color")
+                .unwrap(),
+            "#ABCDEF"
+        );
+
+        // Test retrieving nonexistent color palette
+        assert!(get_palette(&palettes, Some("palette3")).is_none());
+
+        // Test default behavior
+        assert!(get_palette(&palettes, None).is_none());
+    }
 }
diff --git a/src/configs/starship_root.rs b/src/configs/starship_root.rs
index 8edbbad44..aeb03358a 100644
--- a/src/configs/starship_root.rs
+++ b/src/configs/starship_root.rs
@@ -1,4 +1,5 @@
 use serde::{Deserialize, Serialize};
+use std::collections::HashMap;
 
 #[derive(Clone, Serialize, Deserialize, Debug)]
 #[cfg_attr(
@@ -16,8 +17,13 @@ pub struct StarshipRootConfig {
     pub scan_timeout: u64,
     pub command_timeout: u64,
     pub add_newline: bool,
+    #[serde(skip_serializing_if = "Option::is_none")]
+    pub palette: Option<String>,
+    pub palettes: HashMap<String, Palette>,
 }
 
+pub type Palette = HashMap<String, String>;
+
 // List of default prompt order
 // NOTE: If this const value is changed then Default prompt order subheading inside
 // prompt heading of config docs needs to be updated according to changes made here.
@@ -114,6 +120,8 @@ impl Default for StarshipRootConfig {
             scan_timeout: 30,
             command_timeout: 500,
             add_newline: true,
+            palette: None,
+            palettes: HashMap::default(),
         }
     }
 }
diff --git a/src/formatter/string_formatter.rs b/src/formatter/string_formatter.rs
index 8cce8f4e2..2b104acce 100644
--- a/src/formatter/string_formatter.rs
+++ b/src/formatter/string_formatter.rs
@@ -245,7 +245,7 @@ impl<'a> StringFormatter<'a> {
             style_variables: &'a StyleVariableMapType<'a>,
             context: Option<&Context>,
         ) -> Result<Vec<Segment>, StringFormatterError> {
-            let style = parse_style(textgroup.style, style_variables);
+            let style = parse_style(textgroup.style, style_variables, context);
             parse_format(
                 textgroup.format,
                 style.transpose()?,
@@ -258,6 +258,7 @@ impl<'a> StringFormatter<'a> {
         fn parse_style<'a>(
             style: Vec<StyleElement>,
             variables: &'a StyleVariableMapType<'a>,
+            context: Option<&Context>,
         ) -> Option<Result<Style, StringFormatterError>> {
             let style_strings = style
                 .into_iter()
@@ -276,7 +277,7 @@ impl<'a> StringFormatter<'a> {
                 .map(|style_strings| {
                     let style_string: String =
                         style_strings.iter().flat_map(|s| s.chars()).collect();
-                    parse_style_string(&style_string)
+                    parse_style_string(&style_string, context)
                 })
                 .transpose()
         }
diff --git a/src/modules/fill.rs b/src/modules/fill.rs
index bf9558efc..cd5e67c3d 100644
--- a/src/modules/fill.rs
+++ b/src/modules/fill.rs
@@ -14,7 +14,7 @@ pub fn module<'a>(context: &'a Context) -> Option<Module<'a>> {
         return None;
     }
 
-    let style = parse_style_string(config.style);
+    let style = parse_style_string(config.style, Some(context));
 
     module.set_segments(vec![Segment::fill(style, config.symbol)]);