feat(env_var): allow multiple instances (#2797)

Allows displaying multiple instances of the env_var module.

6 files changed, 128 insertions, 23 deletions

diff --git a/docs/config/README.md b/docs/config/README.md
index 2f4c3f379..714ee8e24 100644
--- a/docs/config/README.md
+++ b/docs/config/README.md
@@ -977,12 +977,25 @@ format = "via [ $version](cyan bold) "
 
 ## Environment Variable
 
-The `env_var` module displays the current value of a selected environment variable.
+The `env_var` module displays the current value of a selected environment variables.
 The module will be shown only if any of the following conditions are met:
 
 - The `variable` configuration option matches an existing environment variable
 - The `variable` configuration option is not defined, but the `default` configuration option is
 
+
+::: tip
+Multiple environmental variables can be displayed by using a `.`. (see example)
+If the `variable` configuration option is not set, the module will display value of variable under the name of text after the `.` character.
+
+Example: following configuration will display value of USER environment variable
+```toml
+# ~/.config/starship.toml
+
+[env_var.USER]
+default = "unknown user"
+```
+
 ### Options
 
 | Option     | Default                        | Description                                                                  |
@@ -1013,6 +1026,17 @@ variable = "SHELL"
 default = "unknown shell"
 ```
 
+Displaying multiple environmental variables:
+```toml
+# ~/.config/starship.toml
+
+[env_var.SHELL]
+variable = "SHELL"
+default = "unknown shell"
+[env_var.USER]
+default = "unknown user"
+```
+
 ## Erlang
 
 The `erlang` module shows the currently installed version of [Erlang/OTP](https://erlang.org/doc/).
diff --git a/src/config.rs b/src/config.rs
index 1f002efcb..795f9f38a 100644
--- a/src/config.rs
+++ b/src/config.rs
@@ -342,6 +342,10 @@ impl StarshipConfig {
     pub fn get_custom_modules(&self) -> Option<&toml::value::Table> {
         self.get_config(&["custom"])?.as_table()
     }
+    /// Get the table of all the registered env_var modules, if any
+    pub fn get_env_var_modules(&self) -> Option<&toml::value::Table> {
+        self.get_config(&["env_var"])?.as_table()
+    }
 
     pub fn get_root_config(&self) -> StarshipRootConfig {
         if let Some(root_config) = &self.config {
diff --git a/src/configs/mod.rs b/src/configs/mod.rs
index 658c4553b..f5867dfe4 100644
--- a/src/configs/mod.rs
+++ b/src/configs/mod.rs
@@ -91,7 +91,7 @@ pub struct FullConfig<'a> {
     dotnet: dotnet::DotnetConfig<'a>,
     elixir: elixir::ElixirConfig<'a>,
     elm: elm::ElmConfig<'a>,
-    env_var: env_var::EnvVarConfig<'a>,
+    env_var: IndexMap<String, env_var::EnvVarConfig<'a>>,
     erlang: erlang::ErlangConfig<'a>,
     gcloud: gcloud::GcloudConfig<'a>,
     git_branch: git_branch::GitBranchConfig<'a>,
diff --git a/src/configs/starship_root.rs b/src/configs/starship_root.rs
index 7d5bfcdab..d6dc4617b 100644
--- a/src/configs/starship_root.rs
+++ b/src/configs/starship_root.rs
@@ -105,7 +105,10 @@ impl<'a> ModuleConfig<'a> for StarshipRootConfig<'a> {
                 "command_timeout" => self.command_timeout.load_config(v),
                 "add_newline" => self.add_newline.load_config(v),
                 unknown => {
-                    if !ALL_MODULES.contains(&unknown) && unknown != "custom" {
+                    if !ALL_MODULES.contains(&unknown)
+                        && unknown != "custom"
+                        && unknown != "env_var"
+                    {
                         log::warn!("Unknown config key '{}'", unknown);
 
                         let did_you_mean = &[
@@ -116,6 +119,7 @@ impl<'a> ModuleConfig<'a> for StarshipRootConfig<'a> {
                             "add_newline",
                             // Modules
                             "custom",
+                            "env_var",
                         ]
                         .iter()
                         .chain(ALL_MODULES.iter())
diff --git a/src/modules/env_var.rs b/src/modules/env_var.rs
index 819c992da..846873b58 100644
--- a/src/modules/env_var.rs
+++ b/src/modules/env_var.rs
@@ -3,6 +3,36 @@ use super::{Context, Module};
 use crate::config::RootModuleConfig;
 use crate::configs::env_var::EnvVarConfig;
 use crate::formatter::StringFormatter;
+use crate::segment::Segment;
+
+/// Creates env_var_module displayer which displays all configured environmental variables
+pub fn module<'a>(context: &'a Context) -> Option<Module<'a>> {
+    let config_table = context.config.get_env_var_modules()?;
+    let mut env_modules = config_table
+        .iter()
+        .filter(|(_, config)| config.is_table())
+        .filter_map(|(variable, _)| env_var_module(vec!["env_var", variable], context))
+        .collect::<Vec<Module>>();
+    // Old configuration is present in starship configuration
+    if config_table.iter().any(|(_, config)| !config.is_table()) {
+        if let Some(fallback_env_var_module) = env_var_module(vec!["env_var"], context) {
+            env_modules.push(fallback_env_var_module);
+        }
+    }
+    Some(env_var_displayer(env_modules, context))
+}
+
+/// A utility module to display multiple env_variable modules
+fn env_var_displayer<'a>(modules: Vec<Module>, context: &'a Context) -> Module<'a> {
+    let mut module = context.new_module("env_var_displayer");
+
+    let module_segments = modules
+        .into_iter()
+        .flat_map(|module| module.segments)
+        .collect::<Vec<Segment>>();
+    module.set_segments(module_segments);
+    module
+}
 
 /// Creates a module with the value of the chosen environment variable
 ///
@@ -10,11 +40,20 @@ use crate::formatter::StringFormatter;
 ///     - env_var.disabled is absent or false
 ///     - env_var.variable is defined
 ///     - a variable named as the value of env_var.variable is defined
-pub fn module<'a>(context: &'a Context) -> Option<Module<'a>> {
-    let mut module = context.new_module("env_var");
-    let config: EnvVarConfig = EnvVarConfig::try_load(module.config);
+fn env_var_module<'a>(module_config_path: Vec<&str>, context: &'a Context) -> Option<Module<'a>> {
+    let mut module = context.new_module(&module_config_path.join("."));
+    let config_value = context.config.get_config(&module_config_path);
+    let config = EnvVarConfig::load(config_value.expect(
+        "modules::env_var::module should only be called after ensuring that the module exists",
+    ));
+
+    if config.disabled {
+        return None;
+    };
+
+    let variable_name = get_variable_name(module_config_path, &config);
 
-    let env_value = get_env_value(context, config.variable?, config.default)?;
+    let env_value = get_env_value(context, variable_name?, config.default)?;
     let parsed = StringFormatter::new(config.format).and_then(|formatter| {
         formatter
             .map_meta(|var, _| match var {
@@ -43,6 +82,19 @@ pub fn module<'a>(context: &'a Context) -> Option<Module<'a>> {
     Some(module)
 }
 
+fn get_variable_name<'a>(
+    module_config_path: Vec<&'a str>,
+    config: &'a EnvVarConfig,
+) -> Option<&'a str> {
+    match config.variable {
+        Some(v) => Some(v),
+        None => {
+            let last_element = module_config_path.last()?;
+            Some(*last_element)
+        }
+    }
+}
+
 fn get_env_value(context: &Context, name: &str, default: Option<&str>) -> Option<String> {
     match context.get_env(name) {
         Some(value) => Some(value),
@@ -59,12 +111,22 @@ mod test {
 
     #[test]
     fn empty_config() {
+        let actual = ModuleRenderer::new("env_var").collect();
+        let expected = None;
+
+        assert_eq!(expected, actual);
+    }
+
+    #[test]
+    fn fallback_config() {
         let actual = ModuleRenderer::new("env_var")
             .config(toml::toml! {
                 [env_var]
+                variable="TEST_VAR"
             })
+            .env("TEST_VAR", TEST_VAR_VALUE)
             .collect();
-        let expected = None;
+        let expected = Some(format!("with {} ", style().paint(TEST_VAR_VALUE)));
 
         assert_eq!(expected, actual);
     }
@@ -73,8 +135,7 @@ mod test {
     fn defined_variable() {
         let actual = ModuleRenderer::new("env_var")
             .config(toml::toml! {
-                [env_var]
-                variable = "TEST_VAR"
+                [env_var.TEST_VAR]
             })
             .env("TEST_VAR", TEST_VAR_VALUE)
             .collect();
@@ -87,8 +148,7 @@ mod test {
     fn undefined_variable() {
         let actual = ModuleRenderer::new("env_var")
             .config(toml::toml! {
-                [env_var]
-                variable = "TEST_VAR"
+                [env_var.TEST_VAR]
             })
             .collect();
         let expected = None;
@@ -100,8 +160,7 @@ mod test {
     fn default_has_no_effect() {
         let actual = ModuleRenderer::new("env_var")
             .config(toml::toml! {
-                [env_var]
-                variable = "TEST_VAR"
+                [env_var.TEST_VAR]
                 default = "N/A"
             })
             .env("TEST_VAR", TEST_VAR_VALUE)
@@ -115,8 +174,7 @@ mod test {
     fn default_takes_effect() {
         let actual = ModuleRenderer::new("env_var")
             .config(toml::toml! {
-                [env_var]
-                variable = "UNDEFINED_TEST_VAR"
+                [env_var.UNDEFINED_TEST_VAR]
                 default = "N/A"
             })
             .collect();
@@ -129,8 +187,7 @@ mod test {
     fn symbol() {
         let actual = ModuleRenderer::new("env_var")
             .config(toml::toml! {
-                [env_var]
-                variable = "TEST_VAR"
+                [env_var.TEST_VAR]
                 format = "with [■ $env_value](black bold dimmed) "
             })
             .env("TEST_VAR", TEST_VAR_VALUE)
@@ -147,8 +204,7 @@ mod test {
     fn prefix() {
         let actual = ModuleRenderer::new("env_var")
             .config(toml::toml! {
-                [env_var]
-                variable = "TEST_VAR"
+                [env_var.TEST_VAR]
                 format = "with [_$env_value](black bold dimmed) "
             })
             .env("TEST_VAR", TEST_VAR_VALUE)
@@ -165,8 +221,7 @@ mod test {
     fn suffix() {
         let actual = ModuleRenderer::new("env_var")
             .config(toml::toml! {
-                [env_var]
-                variable = "TEST_VAR"
+                [env_var.TEST_VAR]
                 format = "with [${env_value}_](black bold dimmed) "
             })
             .env("TEST_VAR", TEST_VAR_VALUE)
@@ -179,6 +234,25 @@ mod test {
         assert_eq!(expected, actual);
     }
 
+    #[test]
+    fn display_few() {
+        let actual = ModuleRenderer::new("env_var")
+            .config(toml::toml! {
+                [env_var.TEST_VAR]
+                [env_var.TEST_VAR2]
+            })
+            .env("TEST_VAR", TEST_VAR_VALUE)
+            .env("TEST_VAR2", TEST_VAR_VALUE)
+            .collect();
+        let expected = Some(format!(
+            "with {} with {} ",
+            style().paint(TEST_VAR_VALUE),
+            style().paint(TEST_VAR_VALUE)
+        ));
+
+        assert_eq!(expected, actual);
+    }
+
     fn style() -> Style {
         // default style
         Color::Black.bold().dimmed()
diff --git a/src/modules/mod.rs b/src/modules/mod.rs
index 49c746e70..3742fe13e 100644
--- a/src/modules/mod.rs
+++ b/src/modules/mod.rs
@@ -75,7 +75,6 @@ use std::time::Instant;
 
 pub fn handle<'a>(module: &str, context: &'a Context) -> Option<Module<'a>> {
     let start: Instant = Instant::now();
-
     let mut m: Option<Module> = {
         match module {
             // Keep these ordered alphabetically.