feat: Implement threshold based styling for battery module (#318)

2 files changed, 79 insertions, 34 deletions

diff --git a/src/module.rs b/src/module.rs
index 04233578a..c2f4580e1 100644
--- a/src/module.rs
+++ b/src/module.rs
@@ -145,6 +145,11 @@ impl<'a> Module<'a> {
     pub fn config_value_style(&self, key: &str) -> Option<Style> {
         self.config.and_then(|config| config.get_as_ansi_style(key))
     }
+
+    /// Get a module's config value as an array
+    pub fn config_value_array(&self, key: &str) -> Option<&Vec<toml::Value>> {
+        self.config.and_then(|config| config.get_as_array(key))
+    }
 }
 
 impl<'a> fmt::Display for Module<'a> {
diff --git a/src/modules/battery.rs b/src/modules/battery.rs
index 81edfbf0b..881893940 100644
--- a/src/modules/battery.rs
+++ b/src/modules/battery.rs
@@ -1,13 +1,13 @@
-use ansi_term::Color;
+use ansi_term::{Color, Style};
 
 use super::{Context, Module};
+use crate::config::Config;
 
 /// Creates a module for the battery percentage and charging state
 pub fn module<'a>(context: &'a Context) -> Option<Module<'a>> {
     const BATTERY_FULL: &str = "•";
     const BATTERY_CHARGING: &str = "⇡";
     const BATTERY_DISCHARGING: &str = "⇣";
-    const BATTERY_THRESHOLD: f32 = 10.0;
     // TODO: Update when v1.0 printing refactor is implemented to only
     // print escapes in a prompt context.
     let shell = std::env::var("STARSHIP_SHELL").unwrap_or_default();
@@ -19,43 +19,67 @@ pub fn module<'a>(context: &'a Context) -> Option<Module<'a>> {
     let battery_status = get_battery_status()?;
     let BatteryStatus { state, percentage } = battery_status;
 
-    if percentage > BATTERY_THRESHOLD {
-        log::debug!(
-            "Battery percentage is higher than threshold ({} > {})",
-            percentage,
-            BATTERY_THRESHOLD
-        );
-        return None;
-    }
-
-    // TODO: Set style based on percentage when threshold is modifiable
     let mut module = context.new_module("battery");
-    let module_style = module
-        .config_value_style("style")
-        .unwrap_or_else(|| Color::Red.bold());
-    module.set_style(module_style);
-    module.get_prefix().set_value("");
-
-    match state {
-        battery::State::Full => {
-            module.new_segment("full_symbol", BATTERY_FULL);
-        }
-        battery::State::Charging => {
-            module.new_segment("charging_symbol", BATTERY_CHARGING);
-        }
-        battery::State::Discharging => {
-            module.new_segment("discharging_symbol", BATTERY_DISCHARGING);
+
+    // Parse config under `display`
+    let display_styles = get_display_styles(&module);
+    let display_style = display_styles.iter().find(|display_style| {
+        let BatteryDisplayStyle { threshold, .. } = display_style;
+        percentage <= *threshold as f32
+    });
+
+    if let Some(display_style) = display_style {
+        let BatteryDisplayStyle { style, .. } = display_style;
+
+        // Set style based on percentage
+        module.set_style(*style);
+        module.get_prefix().set_value("");
+
+        match state {
+            battery::State::Full => {
+                module.new_segment("full_symbol", BATTERY_FULL);
+            }
+            battery::State::Charging => {
+                module.new_segment("charging_symbol", BATTERY_CHARGING);
+            }
+            battery::State::Discharging => {
+                module.new_segment("discharging_symbol", BATTERY_DISCHARGING);
+            }
+            _ => return None,
         }
-        _ => return None,
+
+        let mut percent_string = Vec::<String>::with_capacity(2);
+        // Round the percentage to a whole number
+        percent_string.push(percentage.round().to_string());
+        percent_string.push(percentage_char.to_string());
+        module.new_segment("percentage", percent_string.join("").as_ref());
+
+        Some(module)
+    } else {
+        None
     }
+}
 
-    let mut percent_string = Vec::<String>::with_capacity(2);
-    // Round the percentage to a whole number
-    percent_string.push(percentage.round().to_string());
-    percent_string.push(percentage_char.to_string());
-    module.new_segment("percentage", percent_string.join("").as_ref());
+fn get_display_styles(module: &Module) -> Vec<BatteryDisplayStyle> {
+    if let Some(display_configs) = module.config_value_array("display") {
+        let mut display_styles: Vec<BatteryDisplayStyle> = vec![];
+        for display_config in display_configs.iter() {
+            if let toml::Value::Table(config) = display_config {
+                if let Some(display_style) = BatteryDisplayStyle::from_config(config) {
+                    display_styles.push(display_style);
+                }
+            }
+        }
 
-    Some(module)
+        // Return display styles as long as display array exists, even if it is empty.
+        display_styles
+    } else {
+        // Default display styles: [{ threshold = 10, style = "red bold" }]
+        vec![BatteryDisplayStyle {
+            threshold: 10,
+            style: Color::Red.bold(),
+        }]
+    }
 }
 
 fn get_battery_status() -> Option<BatteryStatus> {
@@ -85,3 +109,19 @@ struct BatteryStatus {
     percentage: f32,
     state: battery::State,
 }
+
+#[derive(Clone, Debug)]
+struct BatteryDisplayStyle {
+    threshold: i64,
+    style: Style,
+}
+
+impl BatteryDisplayStyle {
+    /// construct battery display style from toml table
+    pub fn from_config(config: &toml::value::Table) -> Option<BatteryDisplayStyle> {
+        let threshold = config.get_as_i64("threshold")?;
+        let style = config.get_as_ansi_style("style")?;
+
+        Some(BatteryDisplayStyle { threshold, style })
+    }
+}