feat: Allow segment-specific styling (#378)

Adds the ability to style individual segments in the prompt.

The segment documentation is not fully updated in this commit and is waiting on
a config refactor so that we can write unified docs.

2 files changed, 54 insertions, 3 deletions

diff --git a/src/config.rs b/src/config.rs
index 91262658f..70c0754dc 100644
--- a/src/config.rs
+++ b/src/config.rs
@@ -18,6 +18,7 @@ pub trait Config {
     fn get_as_i64(&self, key: &str) -> Option<i64>;
     fn get_as_array(&self, key: &str) -> Option<&Vec<Value>>;
     fn get_as_ansi_style(&self, key: &str) -> Option<ansi_term::Style>;
+    fn get_as_segment_config(&self, key: &str) -> Option<SegmentConfig>;
 
     // Internal implementation for accessors
     fn get_config(&self, key: &str) -> Option<&Value>;
@@ -137,6 +138,39 @@ impl Config for Table {
         self.get_as_str(key)
             .map(|x| parse_style_string(x).unwrap_or_default())
     }
+
+    /// Get a key from a module's configuration as a segment config.
+    ///
+    /// The config can be
+    ///
+    /// - a string, will be interpreted as value.
+    /// - a table with optional { value, style } keys.
+    ///   If omitted, default value will be used.
+    ///
+    /// Returns `Some(SegmentConfig)` if key exists in the configuration, else `None`.
+    fn get_as_segment_config(&self, key: &str) -> Option<SegmentConfig> {
+        self.get_config(key).and_then(|segment_config: &Value| {
+            match segment_config {
+                toml::Value::String(value) => Some(SegmentConfig {
+                    value: Some(value.as_str()),
+                    style: None,
+                }),
+                toml::Value::Table(config_table) => Some(SegmentConfig {
+                    value: config_table.get_as_str("value"),
+                    style: config_table.get_as_ansi_style("style"),
+                }),
+                _ => {
+                    log::debug!(
+                        "Expected \"{}\" to be a string or config table. Instead received {} of type {}.",
+                        key,
+                        segment_config,
+                        segment_config.type_str()
+                        );
+                    None
+                }
+            }
+        })
+    }
 }
 
 fn log_if_key_found(key: &str, something: Option<&Value>) {
@@ -272,6 +306,11 @@ fn parse_color_string(color_string: &str) -> Option<ansi_term::Color> {
     predefined_color
 }
 
+pub struct SegmentConfig<'a> {
+    pub value: Option<&'a str>,
+    pub style: Option<ansi_term::Style>,
+}
+
 #[cfg(test)]
 mod tests {
     use super::*;
diff --git a/src/module.rs b/src/module.rs
index 65e69c36d..b9e71cbec 100644
--- a/src/module.rs
+++ b/src/module.rs
@@ -1,4 +1,5 @@
 use crate::config::Config;
+use crate::config::SegmentConfig;
 use crate::segment::Segment;
 use ansi_term::Style;
 use ansi_term::{ANSIString, ANSIStrings};
@@ -70,9 +71,14 @@ impl<'a> Module<'a> {
     /// Get a reference to a newly created segment in the module
     pub fn new_segment(&mut self, name: &str, value: &str) -> &mut Segment {
         let mut segment = Segment::new(name);
-        segment.set_style(self.style);
-        // Use the provided value unless overwritten by config
-        segment.set_value(self.config_value_str(name).unwrap_or(value));
+        if let Some(segment_config) = self.config_value_segment_config(name) {
+            segment.set_style(segment_config.style.unwrap_or(self.style));
+            segment.set_value(segment_config.value.unwrap_or(value));
+        } else {
+            segment.set_style(self.style);
+            // Use the provided value unless overwritten by config
+            segment.set_value(self.config_value_str(name).unwrap_or(value));
+        }
         self.segments.push(segment);
 
         self.segments.last_mut().unwrap()
@@ -168,6 +174,12 @@ impl<'a> Module<'a> {
     pub fn config_value_array(&self, key: &str) -> Option<&Vec<toml::Value>> {
         self.config.and_then(|config| config.get_as_array(key))
     }
+
+    /// Get a module's config value as a table of segment config
+    pub fn config_value_segment_config(&self, key: &str) -> Option<SegmentConfig> {
+        self.config
+            .and_then(|config| config.get_as_segment_config(key))
+    }
 }
 
 impl<'a> fmt::Display for Module<'a> {