chore: Refactor getting string values from config (#94)

10 files changed, 38 insertions, 27 deletions

diff --git a/src/config.rs b/src/config.rs
index 74ff447d2..4b28d69c5 100644
--- a/src/config.rs
+++ b/src/config.rs
@@ -69,6 +69,7 @@ impl Config {
 /// Extends `toml::value::Table` with useful methods
 pub trait TableExt {
     fn get_as_bool(&self, key: &str) -> Option<bool>;
+    fn get_as_str(&self, key: &str) -> Option<&str>;
 }
 
 impl TableExt for toml::value::Table {
@@ -76,6 +77,11 @@ impl TableExt for toml::value::Table {
     fn get_as_bool(&self, key: &str) -> Option<bool> {
         self.get(key).map(toml::Value::as_bool).unwrap_or(None)
     }
+
+    /// Get a key from a module's configuration as a string
+    fn get_as_str(&self, key: &str) -> Option<&str> {
+        self.get(key).map(toml::Value::as_str).unwrap_or(None)
+    }
 }
 
 mod tests {
@@ -86,14 +92,30 @@ mod tests {
         let mut table = toml::value::Table::new();
 
         // Use with boolean value
-        table.insert("boolean".to_string(), toml::value::Value::Boolean(true));
+        table.insert(String::from("boolean"), toml::value::Value::Boolean(true));
         assert_eq!(table.get_as_bool("boolean"), Some(true));
 
         // Use with string value
         table.insert(
-            "string".to_string(),
-            toml::value::Value::String("true".to_string()),
+            String::from("string"),
+            toml::value::Value::String(String::from("true")),
         );
         assert_eq!(table.get_as_bool("string"), None);
     }
+
+    #[test]
+    fn table_get_as_str() {
+        let mut table = toml::value::Table::new();
+
+        // Use with string value
+        table.insert(
+            String::from("string"),
+            toml::value::Value::String(String::from("hello")),
+        );
+        assert_eq!(table.get_as_str("string"), Some("hello"));
+
+        // Use with boolean value
+        table.insert(String::from("boolean"), toml::value::Value::Boolean(true));
+        assert_eq!(table.get_as_str("boolean"), None);
+    }
 }
diff --git a/src/module.rs b/src/module.rs
index 604258e19..3e530c3f6 100644
--- a/src/module.rs
+++ b/src/module.rs
@@ -1,8 +1,8 @@
+use crate::config::TableExt;
 use crate::segment::Segment;
 use ansi_term::Style;
 use ansi_term::{ANSIString, ANSIStrings};
 use std::fmt;
-use std::string::ToString;
 
 /// A module is a collection of segments showing data for a single integration
 /// (e.g. The git module shows the current git branch and status)
@@ -40,14 +40,11 @@ impl<'a> Module<'a> {
     }
 
     /// Get a reference to a newly created segment in the module
-    pub fn new_segment<T>(&mut self, name: &str, value: T) -> &mut Segment
-    where
-        T: Into<String>,
-    {
+    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(name).unwrap_or_else(|| value.into()));
+        segment.set_value(self.config_value(name).unwrap_or(value));
         self.segments.push(segment);
 
         self.segments.last_mut().unwrap()
@@ -99,16 +96,8 @@ impl<'a> Module<'a> {
     }
 
     /// Get a module's config value as a string
-    fn config_value(&self, key: &str) -> Option<String> {
-        self.config
-            // Find the config value by its key
-            .map(|config| config.get(key))
-            .unwrap_or(None)
-            // Get the config value as a `&str`
-            .map(toml::Value::as_str)
-            .unwrap_or(None)
-            // Convert it to a String
-            .map(str::to_string)
+    fn config_value(&self, key: &str) -> Option<&str> {
+        self.config.and_then(|config| config.get_as_str(key))
     }
 }
 
diff --git a/src/modules/battery.rs b/src/modules/battery.rs
index 048abe05f..586153266 100644
--- a/src/modules/battery.rs
+++ b/src/modules/battery.rs
@@ -43,7 +43,7 @@ pub fn module<'a>(context: &'a Context) -> Option<Module<'a>> {
     // Round the percentage to a whole number
     percent_string.push(percentage.round().to_string());
     percent_string.push("%".to_string());
-    module.new_segment("percentage", percent_string.join(""));
+    module.new_segment("percentage", percent_string.join("").as_ref());
 
     Some(module)
 }
diff --git a/src/modules/directory.rs b/src/modules/directory.rs
index a2cafac16..fe92c5c58 100644
--- a/src/modules/directory.rs
+++ b/src/modules/directory.rs
@@ -38,7 +38,7 @@ pub fn module<'a>(context: &'a Context) -> Option<Module<'a>> {
 
     // Truncate the dir string to the maximum number of path components
     let truncated_dir_string = truncate(dir_string, DIR_TRUNCATION_LENGTH);
-    module.new_segment("path", truncated_dir_string);
+    module.new_segment("path", &truncated_dir_string);
 
     module.get_prefix().set_value("in ");
 
diff --git a/src/modules/git_branch.rs b/src/modules/git_branch.rs
index d2861c47b..06aa6320d 100644
--- a/src/modules/git_branch.rs
+++ b/src/modules/git_branch.rs
@@ -16,7 +16,7 @@ pub fn module<'a>(context: &'a Context) -> Option<Module<'a>> {
     module.get_prefix().set_value("on ");
 
     module.new_segment("branch_char", GIT_BRANCH_CHAR);
-    module.new_segment("branch_name", branch_name.to_string());
+    module.new_segment("branch_name", branch_name);
 
     Some(module)
 }
diff --git a/src/modules/go.rs b/src/modules/go.rs
index a1b192e4a..a50356d50 100644
--- a/src/modules/go.rs
+++ b/src/modules/go.rs
@@ -35,7 +35,7 @@ pub fn module<'a>(context: &'a Context) -> Option<Module<'a>> {
 
             let formatted_version = format_go_version(go_version)?;
             module.new_segment("symbol", GO_CHAR);
-            module.new_segment("version", formatted_version);
+            module.new_segment("version", &formatted_version);
 
             Some(module)
         }
diff --git a/src/modules/package.rs b/src/modules/package.rs
index 527d5eaec..03d2b6d53 100644
--- a/src/modules/package.rs
+++ b/src/modules/package.rs
@@ -19,7 +19,7 @@ pub fn module<'a>(context: &'a Context) -> Option<Module<'a>> {
             module.get_prefix().set_value("is ");
 
             module.new_segment("symbol", PACKAGE_CHAR);
-            module.new_segment("version", package_version);
+            module.new_segment("version", &package_version);
 
             Some(module)
         }
diff --git a/src/modules/python.rs b/src/modules/python.rs
index a1a36cd3a..dbee151c8 100644
--- a/src/modules/python.rs
+++ b/src/modules/python.rs
@@ -31,7 +31,7 @@ pub fn module<'a>(context: &'a Context) -> Option<Module<'a>> {
 
             let formatted_version = format_python_version(python_version);
             module.new_segment("symbol", PYTHON_CHAR);
-            module.new_segment("version", formatted_version);
+            module.new_segment("version", &formatted_version);
 
             Some(module)
         }
diff --git a/src/modules/rust.rs b/src/modules/rust.rs
index 057344064..a89ae39d1 100644
--- a/src/modules/rust.rs
+++ b/src/modules/rust.rs
@@ -29,7 +29,7 @@ pub fn module<'a>(context: &'a Context) -> Option<Module<'a>> {
 
             let formatted_version = format_rustc_version(rust_version);
             module.new_segment("symbol", RUST_CHAR);
-            module.new_segment("version", formatted_version);
+            module.new_segment("version", &formatted_version);
 
             Some(module)
         }
diff --git a/src/modules/username.rs b/src/modules/username.rs
index 73295b79e..745823fcb 100644
--- a/src/modules/username.rs
+++ b/src/modules/username.rs
@@ -20,7 +20,7 @@ pub fn module<'a>(context: &'a Context) -> Option<Module<'a>> {
     if user != logname || ssh_connection.is_some() || is_root(&mut module_color) {
         let mut module = context.new_module("username")?;
         module.set_style(module_color);
-        module.new_segment("username", user?);
+        module.new_segment("username", &user?);
 
         return Some(module);
     }