feat(env_var): Add support for env_var.VAR in format (#4497)

Co-Authored-By: Segev Finer <24731903+segevfiner@users.noreply.github.com>

Co-authored-by: Segev Finer <24731903+segevfiner@users.noreply.github.com>

.github/config-schema.json

@@ -2834,6 +2834,10 @@
         "disabled": {
           "default": false,
           "type": "boolean"
+        },
+        "description": {
+          "default": "<env_var module>",
+          "type": "string"
         }
       },
       "additionalProperties": false

docs/config/README.md

@@ -1380,6 +1380,14 @@ The module will be shown only if any of the following conditions are met:
 
 ::: tip
 
+The order in which env_var modules are shown can be individually set by including
+`${env_var.foo}` in the top level `format` (as it includes a dot, you need to use `${...}`).
+By default, the `env_var` module will simply show all env_var modules in the order they were defined.
+
+:::
+
+::: 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.
 
@@ -1396,13 +1404,14 @@ default = 'unknown user'
 
 ### Options
 
-| Option     | Default                        | Description                                                                  |
+| Option        | Default                        | Description                                                                  |
-| ---------- | ------------------------------ | ---------------------------------------------------------------------------- |
+| ------------- | ------------------------------ | ---------------------------------------------------------------------------- |
-| `symbol`   | `''`                           | The symbol used before displaying the variable value.                        |
+| `symbol`      | `""`                           | The symbol used before displaying the variable value.                        |
-| `variable` |                                | The environment variable to be displayed.                                    |
+| `variable`    |                                | The environment variable to be displayed.                                    |
-| `default`  |                                | The default value to be displayed when the selected variable is not defined. |
+| `default`     |                                | The default value to be displayed when the selected variable is not defined. |
-| `format`   | `'with [$env_value]($style) '` | The format for the module.                                                   |
+| `format`      | `"with [$env_value]($style) "` | The format for the module.                                                   |
-| `disabled` | `false`                        | Disables the `env_var` module.                                               |
+| `description` | `"<env_var module>"`           | The description of the module that is shown when running `starship explain`. |
+| `disabled`    | `false`                        | Disables the `env_var` module.                                               |
 
 ### Variables
 

src/configs/env_var.rs

@@ -16,6 +16,7 @@ pub struct EnvVarConfig<'a> {
     pub default: Option<&'a str>,
     pub format: &'a str,
     pub disabled: bool,
+    pub description: &'a str,
 }
 
 impl<'a> Default for EnvVarConfig<'a> {
@@ -27,6 +28,7 @@ impl<'a> Default for EnvVarConfig<'a> {
             default: None,
             format: "with [$env_value]($style) ",
             disabled: false,
+            description: "<env_var module>",
         }
     }
 }

src/context.rs

@@ -238,15 +238,6 @@ impl<'a> Context<'a> {
         disabled == Some(true)
     }
 
-    /// Return whether the specified custom module has a `disabled` option set to true.
-    /// If it doesn't exist, `None` is returned.
-    pub fn is_custom_module_disabled_in_config(&self, name: &str) -> Option<bool> {
-        let config = self.config.get_custom_module_config(name)?;
-        let disabled = Some(config).and_then(|table| table.as_table()?.get("disabled")?.as_bool());
-
-        Some(disabled == Some(true))
-    }
-
     // returns a new ScanDir struct with reference to current dir_files of context
     // see ScanDir for methods
     pub fn try_begin_scan(&'a self) -> Option<ScanDir<'a>> {

src/module.rs

@@ -31,7 +31,6 @@ pub const ALL_MODULES: &[&str] = &[
     "dotnet",
     "elixir",
     "elm",
-    "env_var",
     "erlang",
     "fennel",
     "fill",

src/modules/custom.rs

@@ -1,9 +1,9 @@
 use std::env;
+use std::fmt::{self, Debug};
 use std::io::Write;
 use std::path::Path;
 use std::process::{Command, Stdio};
 use std::time::Duration;
-use std::time::Instant;
 
 use process_control::{ChildExt, Control, Output};
 
@@ -22,11 +22,11 @@ use crate::{
 ///
 /// Finally, the content of the module itself is also set by a command.
 pub fn module<'a>(name: &str, context: &'a Context) -> Option<Module<'a>> {
-    let start: Instant = Instant::now();
+    let toml_config = get_config(name, context)?;
-    let toml_config = context.config.get_custom_module_config(name).expect(
-        "modules::custom::module should only be called after ensuring that the module exists",
-    );
     let config = CustomConfig::load(toml_config);
+    if config.disabled {
+        return None;
+    }
 
     if let Some(os) = config.os {
         if os != env::consts::OS && !(os == "unix" && cfg!(unix)) {
@@ -34,7 +34,8 @@ pub fn module<'a>(name: &str, context: &'a Context) -> Option<Module<'a>> {
         }
     }
 
-    let mut module = Module::new(name, config.description, Some(toml_config));
+    // Note: Forward config if `Module` ends up needing `config`
+    let mut module = Module::new(&format!("custom.{name}"), config.description, None);
 
     let mut is_match = context
         .try_begin_scan()?
@@ -48,48 +49,74 @@ pub fn module<'a>(name: &str, context: &'a Context) -> Option<Module<'a>> {
             Either::First(b) => b,
             Either::Second(s) => exec_when(s, &config, context),
         };
+
+        if !is_match {
+            return None;
+        }
     }
 
-    if is_match {
+    let parsed = StringFormatter::new(config.format).and_then(|formatter| {
-        let parsed = StringFormatter::new(config.format).and_then(|formatter| {
+        formatter
-            formatter
+            .map_meta(|var, _| match var {
-                .map_meta(|var, _| match var {
+                "symbol" => Some(config.symbol),
-                    "symbol" => Some(config.symbol),
+                _ => None,
-                    _ => None,
+            })
-                })
+            .map_style(|variable| match variable {
-                .map_style(|variable| match variable {
+                "style" => Some(Ok(config.style)),
-                    "style" => Some(Ok(config.style)),
+                _ => None,
-                    _ => None,
+            })
-                })
+            .map_no_escaping(|variable| match variable {
-                .map_no_escaping(|variable| match variable {
+                "output" => {
-                    "output" => {
+                    let output = exec_command(config.command, context, &config)?;
-                        let output = exec_command(config.command, context, &config)?;
+                    let trimmed = output.trim();
-                        let trimmed = output.trim();
 
-                        if trimmed.is_empty() {
+                    if trimmed.is_empty() {
-                            None
+                        None
-                        } else {
+                    } else {
-                            Some(Ok(trimmed.to_string()))
+                        Some(Ok(trimmed.to_string()))
-                        }
                     }
-                    _ => None,
+                }
-                })
+                _ => None,
-                .parse(None, Some(context))
+            })
-        });
+            .parse(None, Some(context))
+    });
 
-        match parsed {
+    match parsed {
-            Ok(segments) => module.set_segments(segments),
+        Ok(segments) => module.set_segments(segments),
-            Err(error) => {
+        Err(error) => {
-                log::warn!("Error in module `custom.{}`:\n{}", name, error);
+            log::warn!("Error in module `custom.{}`:\n{}", name, error);
-            }
+        }
-        };
+    };
-    }
-    let elapsed = start.elapsed();
-    log::trace!("Took {:?} to compute custom module {:?}", elapsed, name);
-    module.duration = elapsed;
     Some(module)
 }
 
+/// Gets the TOML config for the custom module, handling the case where the module is not defined
+fn get_config<'a>(module_name: &str, context: &'a Context<'a>) -> Option<&'a toml::Value> {
+    struct DebugCustomModules<'tmp>(&'tmp toml::value::Table);
+
+    impl Debug for DebugCustomModules<'_> {
+        fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
+            f.debug_list().entries(self.0.keys()).finish()
+        }
+    }
+
+    let config = context.config.get_custom_module_config(module_name);
+
+    if config.is_some() {
+        return config;
+    } else if let Some(modules) = context.config.get_custom_modules() {
+        log::debug!(
+                "top level format contains custom module {module_name:?}, but no configuration was provided. Configuration for the following modules were provided: {:?}",
+                DebugCustomModules(modules),
+        );
+    } else {
+        log::debug!(
+            "top level format contains custom module {module_name:?}, but no configuration was provided.",
+        );
+    };
+    None
+}
+
 /// Return the invoking shell, using `shell` and fallbacking in order to `STARSHIP_SHELL` and "sh"/"cmd"
 fn get_shell<'a, 'b>(
     shell_args: &'b [&'a str],
@@ -680,4 +707,18 @@ mod tests {
 
         dir.close()
     }
+
+    #[test]
+    fn disabled() {
+        let actual = ModuleRenderer::new("custom.test")
+            .config(toml::toml! {
+                [custom.test]
+                disabled = true
+                when = true
+                format = "test"
+            })
+            .collect();
+        let expected = None;
+        assert_eq!(expected, actual);
+    }
 }

src/modules/env_var.rs

@@ -1,38 +1,9 @@
 use super::{Context, Module};
+use std::borrow::Cow;
 
 use crate::config::ModuleConfig;
 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
 ///
@@ -40,20 +11,36 @@ fn env_var_displayer<'a>(modules: Vec<Module>, context: &'a Context) -> Module<'
 ///     - `env_var.disabled` is absent or false
 ///     - `env_var.variable` is defined
 ///     - a variable named as the value of `env_var.variable` is defined
-fn env_var_module<'a>(module_config_path: Vec<&str>, context: &'a Context) -> Option<Module<'a>> {
+pub fn module<'a>(name: Option<&str>, context: &'a Context) -> Option<Module<'a>> {
-    let mut module = context.new_module(&module_config_path.join("."));
+    let toml_config = match name {
-    let config_value = context.config.get_config(&module_config_path);
+        Some(name) => context
-    let config = EnvVarConfig::load(config_value.expect(
+            .config
-        "modules::env_var::module should only be called after ensuring that the module exists",
+            .get_config(&["env_var", name])
-    ));
+            .map(Cow::Borrowed),
+        None => context
+            .config
+            .get_module_config("env_var")
+            .and_then(filter_config)
+            .map(Cow::Owned)
+            .map(Some)?,
+    };
 
+    let mod_name = match name {
+        Some(name) => format!("env_var.{}", name),
+        None => "env_var".to_owned(),
+    };
+
+    let config = EnvVarConfig::try_load(toml_config.as_deref());
+    // Note: Forward config if `Module` ends up needing `config`
+    let mut module = Module::new(&mod_name, config.description, None);
     if config.disabled {
         return None;
     };
 
-    let variable_name = get_variable_name(module_config_path, &config);
+    let variable_name = config.variable.or(name)?;
 
-    let env_value = get_env_value(context, variable_name?, config.default)?;
+    let env_value = context.get_env(variable_name);
+    let env_value = env_value.as_deref().or(config.default)?;
     let parsed = StringFormatter::new(config.format).and_then(|formatter| {
         formatter
             .map_meta(|var, _| match var {
@@ -65,7 +52,7 @@ fn env_var_module<'a>(module_config_path: Vec<&str>, context: &'a Context) -> Op
                 _ => None,
             })
             .map(|variable| match variable {
-                "env_value" => Some(Ok(&env_value)),
+                "env_value" => Some(Ok(env_value)),
                 _ => None,
             })
             .parse(None, Some(context))
@@ -82,24 +69,22 @@ fn env_var_module<'a>(module_config_path: Vec<&str>, context: &'a Context) -> Op
     Some(module)
 }
 
-fn get_variable_name<'a>(
+/// Filter `config` to only includes non-table values
-    module_config_path: Vec<&'a str>,
+/// This filters the top-level table to only include its specific configuation
-    config: &'a EnvVarConfig,
+fn filter_config(config: &toml::Value) -> Option<toml::Value> {
-) -> Option<&'a str> {
+    let o = config
-    match config.variable {
+        .as_table()
-        Some(v) => Some(v),
+        .map(|table| {
-        None => {
+            table
-            let last_element = module_config_path.last()?;
+                .iter()
-            Some(*last_element)
+                .filter(|(_key, val)| !val.is_table())
-        }
+                .map(|(key, val)| (key.to_owned(), val.to_owned()))
-    }
+                .collect::<toml::value::Table>()
-}
+        })
-
+        .filter(|table| !table.is_empty())
-fn get_env_value(context: &Context, name: &str, default: Option<&str>) -> Option<String> {
+        .map(toml::Value::Table);
-    match context.get_env(name) {
+    log::trace!("Filtered top-level env_var config: {o:?}");
-        Some(value) => Some(value),
+    o
-        None => default.map(std::borrow::ToOwned::to_owned),
-    }
 }
 
 #[cfg(test)]
@@ -133,7 +118,7 @@ mod test {
 
     #[test]
     fn defined_variable() {
-        let actual = ModuleRenderer::new("env_var")
+        let actual = ModuleRenderer::new("env_var.TEST_VAR")
             .config(toml::toml! {
                 [env_var.TEST_VAR]
             })
@@ -146,7 +131,7 @@ mod test {
 
     #[test]
     fn undefined_variable() {
-        let actual = ModuleRenderer::new("env_var")
+        let actual = ModuleRenderer::new("env_var.TEST_VAR")
             .config(toml::toml! {
                 [env_var.TEST_VAR]
             })
@@ -158,7 +143,7 @@ mod test {
 
     #[test]
     fn default_has_no_effect() {
-        let actual = ModuleRenderer::new("env_var")
+        let actual = ModuleRenderer::new("env_var.TEST_VAR")
             .config(toml::toml! {
                 [env_var.TEST_VAR]
                 default = "N/A"
@@ -172,7 +157,7 @@ mod test {
 
     #[test]
     fn default_takes_effect() {
-        let actual = ModuleRenderer::new("env_var")
+        let actual = ModuleRenderer::new("env_var.UNDEFINED_TEST_VAR")
             .config(toml::toml! {
                 [env_var.UNDEFINED_TEST_VAR]
                 default = "N/A"
@@ -185,7 +170,7 @@ mod test {
 
     #[test]
     fn symbol() {
-        let actual = ModuleRenderer::new("env_var")
+        let actual = ModuleRenderer::new("env_var.TEST_VAR")
             .config(toml::toml! {
                 [env_var.TEST_VAR]
                 format = "with [■ $env_value](black bold dimmed) "
@@ -202,7 +187,7 @@ mod test {
 
     #[test]
     fn prefix() {
-        let actual = ModuleRenderer::new("env_var")
+        let actual = ModuleRenderer::new("env_var.TEST_VAR")
             .config(toml::toml! {
                 [env_var.TEST_VAR]
                 format = "with [_$env_value](black bold dimmed) "
@@ -219,7 +204,7 @@ mod test {
 
     #[test]
     fn suffix() {
-        let actual = ModuleRenderer::new("env_var")
+        let actual = ModuleRenderer::new("env_var.TEST_VAR")
             .config(toml::toml! {
                 [env_var.TEST_VAR]
                 format = "with [${env_value}_](black bold dimmed) "
@@ -236,7 +221,7 @@ mod test {
 
     #[test]
     fn display_few() {
-        let actual = ModuleRenderer::new("env_var")
+        let actual1 = ModuleRenderer::new("env_var.TEST_VAR")
             .config(toml::toml! {
                 [env_var.TEST_VAR]
                 [env_var.TEST_VAR2]
@@ -244,11 +229,103 @@ mod test {
             .env("TEST_VAR", TEST_VAR_VALUE)
             .env("TEST_VAR2", TEST_VAR_VALUE)
             .collect();
-        let expected = Some(format!(
+        let actual2 = ModuleRenderer::new("env_var.TEST_VAR2")
-            "with {} with {} ",
+            .config(toml::toml! {
-            style().paint(TEST_VAR_VALUE),
+                [env_var.TEST_VAR]
-            style().paint(TEST_VAR_VALUE)
+                [env_var.TEST_VAR2]
-        ));
+            })
+            .env("TEST_VAR", TEST_VAR_VALUE)
+            .env("TEST_VAR2", TEST_VAR_VALUE)
+            .collect();
+        let expected = Some(format!("with {} ", style().paint(TEST_VAR_VALUE)));
+
+        assert_eq!(expected, actual1);
+        assert_eq!(expected, actual2);
+    }
+
+    #[test]
+    fn mixed() {
+        let cfg = toml::toml! {
+            [env_var]
+            variable = "TEST_VAR_OUTER"
+            format = "$env_value"
+            [env_var.TEST_VAR_INNER]
+            format = "$env_value"
+        };
+        let actual_inner = ModuleRenderer::new("env_var.TEST_VAR_INNER")
+            .config(cfg.clone())
+            .env("TEST_VAR_OUTER", "outer")
+            .env("TEST_VAR_INNER", "inner")
+            .collect();
+
+        assert_eq!(
+            actual_inner.as_deref(),
+            Some("inner"),
+            "inner module should be rendered"
+        );
+
+        let actual_outer = ModuleRenderer::new("env_var")
+            .config(cfg)
+            .env("TEST_VAR_OUTER", "outer")
+            .env("TEST_VAR_INNER", "inner")
+            .collect();
+
+        assert_eq!(
+            actual_outer.as_deref(),
+            Some("outer"),
+            "outer module should be rendered"
+        );
+    }
+
+    #[test]
+    fn no_config() {
+        let actual = ModuleRenderer::new("env_var.TEST_VAR")
+            .env("TEST_VAR", TEST_VAR_VALUE)
+            .collect();
+        let expected = Some(format!("with {} ", style().paint(TEST_VAR_VALUE)));
+
+        assert_eq!(expected, actual);
+    }
+
+    #[test]
+    fn disabled_child() {
+        let actual = ModuleRenderer::new("env_var.TEST_VAR")
+            .config(toml::toml! {
+                [env_var.TEST_VAR]
+                disabled = true
+            })
+            .env("TEST_VAR", TEST_VAR_VALUE)
+            .collect();
+        let expected = None;
+
+        assert_eq!(expected, actual);
+    }
+
+    #[test]
+    fn disabled_root() {
+        let actual = ModuleRenderer::new("env_var")
+            .config(toml::toml! {
+                [env_var]
+                disabled = true
+            })
+            .env("TEST_VAR", TEST_VAR_VALUE)
+            .collect();
+        let expected = None;
+
+        assert_eq!(expected, actual);
+    }
+
+    #[test]
+    fn variable_override() {
+        let actual = ModuleRenderer::new("env_var.TEST_VAR")
+            .config(toml::toml! {
+                [env_var.TEST_VAR]
+                variable = "TEST_VAR2"
+            })
+            .env("TEST_VAR", "implicit name")
+            .env("TEST_VAR2", "explicit name")
+            .collect();
+        let expected = Some(format!("with {} ", style().paint("explicit name")));
 
         assert_eq!(expected, actual);
     }

src/modules/mod.rs

@@ -122,8 +122,8 @@ pub fn handle<'a>(module: &str, context: &'a Context) -> Option<Module<'a>> {
             "elixir" => elixir::module(context),
             "elm" => elm::module(context),
             "erlang" => erlang::module(context),
+            "env_var" => env_var::module(None, context),
             "fennel" => fennel::module(context),
-            "env_var" => env_var::module(context),
             "fill" => fill::module(context),
             "gcloud" => gcloud::module(context),
             "git_branch" => git_branch::module(context),
@@ -183,8 +183,9 @@ pub fn handle<'a>(module: &str, context: &'a Context) -> Option<Module<'a>> {
             "vagrant" => vagrant::module(context),
             "vcsh" => vcsh::module(context),
             "zig" => zig::module(context),
-            // Added for tests, avoid potential side effects in production code.
+            env if env.starts_with("env_var.") => {
-            #[cfg(test)]
+                env_var::module(env.strip_prefix("env_var."), context)
+            }
             custom if custom.starts_with("custom.") => {
                 // SAFETY: We just checked that the module starts with "custom."
                 custom::module(custom.strip_prefix("custom.").unwrap(), context)
@@ -233,7 +234,6 @@ pub fn description(module: &str) -> &'static str {
         "dotnet" => "The relevant version of the .NET Core SDK for the current directory",
         "elixir" => "The currently installed versions of Elixir and OTP",
         "elm" => "The currently installed version of Elm",
-        "env_var" => "Displays the current value of a selected environment variable",
         "erlang" => "Current OTP version",
         "fennel" => "The currently installed version of Fennel",
         "fill" => "Fills the remaining space on the line with a pad string",

src/print.rs

@@ -2,7 +2,7 @@ use clap::{builder::PossibleValue, ValueEnum};
 use nu_ansi_term::AnsiStrings;
 use rayon::prelude::*;
 use std::collections::BTreeSet;
-use std::fmt::{self, Debug, Write as FmtWrite};
+use std::fmt::{Debug, Write as FmtWrite};
 use std::io::{self, Write};
 use std::time::Duration;
 use terminal_size::terminal_size;
@@ -313,14 +313,6 @@ fn handle_module<'a>(
     context: &'a Context,
     module_list: &BTreeSet<String>,
 ) -> Vec<Module<'a>> {
-    struct DebugCustomModules<'tmp>(&'tmp toml::value::Table);
-
-    impl Debug for DebugCustomModules<'_> {
-        fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
-            f.debug_list().entries(self.0.keys()).finish()
-        }
-    }
-
     let mut modules: Vec<Module> = Vec::new();
 
     if ALL_MODULES.contains(&module) {
@@ -328,34 +320,29 @@ fn handle_module<'a>(
         if !context.is_module_disabled_in_config(module) {
             modules.extend(modules::handle(module, context));
         }
-    } else if module == "custom" {
+    } else if module.starts_with("custom.") || module.starts_with("env_var.") {
-        // Write out all custom modules, except for those that are explicitly set
+        // custom.<name> and env_var.<name> are special cases and handle disabled modules themselves
-        if let Some(custom_modules) = context.config.get_custom_modules() {
+        modules.extend(modules::handle(module, context));
-            let custom_modules = custom_modules.iter().filter_map(|(custom_module, config)| {
+    } else if matches!(module, "custom" | "env_var") {
-                if should_add_implicit_custom_module(custom_module, config, module_list) {
+        // env var is a spacial case and may contain a top-level module definition
-                    modules::custom::module(custom_module, context)
+        if module == "env_var" {
-                } else {
+            modules.extend(modules::handle(module, context));
-                    None
-                }
-            });
-            modules.extend(custom_modules);
         }
-    } else if let Some(module) = module.strip_prefix("custom.") {
+
-        // Write out a custom module if it isn't disabled (and it exists...)
+        // Write out all custom modules, except for those that are explicitly set
-        match context.is_custom_module_disabled_in_config(module) {
+        for (child, config) in context
-            Some(true) => (), // Module is disabled, we don't add it to the prompt
+            .config
-            Some(false) => modules.extend(modules::custom::module(module, context)),
+            .get_config(&[module])
-            None => match context.config.get_custom_modules() {
+            .and_then(|config| config.as_table().map(|t| t.iter()))
-                Some(modules) => log::debug!(
+            .into_iter()
-                    "top level format contains custom module \"{}\", but no configuration was provided. Configuration for the following modules were provided: {:?}",
+            .flatten()
-                    module,
+        {
-                    DebugCustomModules(modules),
+            // Some env var keys may be part of a top-level module definition
-                    ),
+            if module == "env_var" && !config.is_table() {
-                None => log::debug!(
+                continue;
-                    "top level format contains custom module \"{}\", but no configuration was provided.",
+            } else if should_add_implicit_module(module, child, config, module_list) {
-                    module,
+                modules.extend(modules::handle(&format!("{module}.{child}"), context));
-                    ),
+            }
-            },
         }
     } else {
         log::debug!(
@@ -368,12 +355,13 @@ fn handle_module<'a>(
     modules
 }
 
-fn should_add_implicit_custom_module(
+fn should_add_implicit_module(
-    custom_module: &str,
+    parent_module: &str,
+    child_module: &str,
     config: &toml::Value,
     module_list: &BTreeSet<String>,
 ) -> bool {
-    let explicit_module_name = format!("custom.{custom_module}");
+    let explicit_module_name = format!("{parent_module}.{child_module}");
     let is_explicitly_specified = module_list.contains(&explicit_module_name);
 
     if is_explicitly_specified {
@@ -539,4 +527,153 @@ mod test {
     fn print_schema_does_not_panic() {
         print_schema();
     }
+
+    #[test]
+    fn custom_expands() -> std::io::Result<()> {
+        let dir = tempfile::tempdir()?;
+        let mut context = default_context();
+        context.current_dir = dir.path().to_path_buf();
+        context.config = StarshipConfig {
+            config: Some(toml::toml! {
+                format="$custom"
+                [custom.a]
+                when=true
+                format="a"
+                [custom.b]
+                when=true
+                format="b"
+            }),
+        };
+        context.root_config.format = "$custom".to_string();
+
+        let expected = String::from("\nab");
+        let actual = get_prompt(context);
+        assert_eq!(expected, actual);
+        dir.close()
+    }
+
+    #[test]
+    fn env_expands() {
+        let mut context = default_context();
+        context.config = StarshipConfig {
+            config: Some(toml::toml! {
+                format="$env_var"
+                [env_var]
+                format="$env_value"
+                variable = "a"
+                [env_var.b]
+                format="$env_value"
+                [env_var.c]
+                format="$env_value"
+            }),
+        };
+        context.root_config.format = "$env_var".to_string();
+        context.env.insert("a", "a".to_string());
+        context.env.insert("b", "b".to_string());
+        context.env.insert("c", "c".to_string());
+
+        let expected = String::from("\nabc");
+        let actual = get_prompt(context);
+        assert_eq!(expected, actual);
+    }
+
+    #[test]
+    fn custom_mixed() -> std::io::Result<()> {
+        let dir = tempfile::tempdir()?;
+        let mut context = default_context();
+        context.current_dir = dir.path().to_path_buf();
+        context.config = StarshipConfig {
+            config: Some(toml::toml! {
+                format="${custom.c}$custom${custom.b}"
+                [custom.a]
+                when=true
+                format="a"
+                [custom.b]
+                when=true
+                format="b"
+                [custom.c]
+                when=true
+                format="c"
+            }),
+        };
+        context.root_config.format = "${custom.c}$custom${custom.b}".to_string();
+
+        let expected = String::from("\ncab");
+        let actual = get_prompt(context);
+        assert_eq!(expected, actual);
+        dir.close()
+    }
+
+    #[test]
+    fn env_mixed() {
+        let mut context = default_context();
+        context.config = StarshipConfig {
+            config: Some(toml::toml! {
+                format="${env_var.c}$env_var${env_var.b}"
+                [env_var]
+                format="$env_value"
+                variable = "d"
+                [env_var.a]
+                format="$env_value"
+                [env_var.b]
+                format="$env_value"
+                [env_var.c]
+                format="$env_value"
+            }),
+        };
+        context.root_config.format = "${env_var.c}$env_var${env_var.b}".to_string();
+        context.env.insert("a", "a".to_string());
+        context.env.insert("b", "b".to_string());
+        context.env.insert("c", "c".to_string());
+        context.env.insert("d", "d".to_string());
+
+        let expected = String::from("\ncdab");
+        let actual = get_prompt(context);
+        assert_eq!(expected, actual);
+    }
+
+    #[test]
+    fn custom_subset() -> std::io::Result<()> {
+        let dir = tempfile::tempdir()?;
+        let mut context = default_context();
+        context.current_dir = dir.path().to_path_buf();
+        context.config = StarshipConfig {
+            config: Some(toml::toml! {
+                format="${custom.b}"
+                [custom.a]
+                when=true
+                format="a"
+                [custom.b]
+                when=true
+                format="b"
+            }),
+        };
+        context.root_config.format = "${custom.b}".to_string();
+
+        let expected = String::from("\nb");
+        let actual = get_prompt(context);
+        assert_eq!(expected, actual);
+        dir.close()
+    }
+
+    #[test]
+    fn custom_missing() -> std::io::Result<()> {
+        let dir = tempfile::tempdir()?;
+        let mut context = default_context();
+        context.current_dir = dir.path().to_path_buf();
+        context.config = StarshipConfig {
+            config: Some(toml::toml! {
+                format="${custom.b}"
+                [custom.a]
+                when=true
+                format="a"
+            }),
+        };
+        context.root_config.format = "${custom.b}".to_string();
+
+        let expected = String::from("\n");
+        let actual = get_prompt(context);
+        assert_eq!(expected, actual);
+        dir.close()
+    }
 }