feat: refactor modules to use format strings (#1374)

1 files changed, 118 insertions, 80 deletions

diff --git a/src/print.rs b/src/print.rs
index 369f3b387..5491a694e 100644
--- a/src/print.rs
+++ b/src/print.rs
@@ -1,14 +1,18 @@
 use ansi_term::ANSIStrings;
 use clap::ArgMatches;
 use rayon::prelude::*;
+use std::collections::BTreeSet;
 use std::fmt::{self, Debug, Write as FmtWrite};
 use std::io::{self, Write};
 use unicode_width::UnicodeWidthChar;
 
+use crate::configs::PROMPT_ORDER;
 use crate::context::{Context, Shell};
+use crate::formatter::{StringFormatter, VariableHolder};
 use crate::module::Module;
 use crate::module::ALL_MODULES;
 use crate::modules;
+use crate::segment::Segment;
 
 pub fn prompt(args: ArgMatches) {
     let context = Context::new(args);
@@ -21,34 +25,53 @@ pub fn get_prompt(context: Context) -> String {
     let config = context.config.get_root_config();
     let mut buf = String::new();
 
-    // Write a new line before the prompt
-    if config.add_newline {
-        writeln!(buf).unwrap();
-    }
-
     // A workaround for a fish bug (see #739,#279). Applying it to all shells
     // breaks things (see #808,#824,#834). Should only be printed in fish.
     if let Shell::Fish = context.shell {
         buf.push_str("\x1b[J"); // An ASCII control code to clear screen
     }
 
-    let modules = compute_modules(&context);
-
-    let mut print_without_prefix = true;
-    let printable = modules.iter();
-
-    for module in printable {
-        // Skip printing the prefix of a module after the line_break
-        if print_without_prefix {
-            let module_without_prefix = module.to_string_without_prefix(context.shell);
-            write!(buf, "{}", module_without_prefix).unwrap()
+    let formatter = if let Ok(formatter) = StringFormatter::new(config.format) {
+        formatter
+    } else {
+        log::error!("Error parsing `format`");
+        buf.push_str(">");
+        return buf;
+    };
+    let modules = formatter.get_variables();
+    let formatter = formatter.map_variables_to_segments(|module| {
+        // Make $all display all modules
+        if module == "all" {
+            Some(Ok(PROMPT_ORDER
+                .par_iter()
+                .flat_map(|module| {
+                    handle_module(module, &context, &modules)
+                        .into_iter()
+                        .flat_map(|module| module.segments)
+                        .collect::<Vec<Segment>>()
+                })
+                .collect::<Vec<_>>()))
+        } else if context.is_module_disabled_in_config(&module) {
+            None
         } else {
-            let module = module.ansi_strings_for_shell(context.shell);
-            write!(buf, "{}", ANSIStrings(&module)).unwrap();
+            // Get segments from module
+            Some(Ok(handle_module(module, &context, &modules)
+                .into_iter()
+                .flat_map(|module| module.segments)
+                .collect::<Vec<Segment>>()))
         }
+    });
 
-        print_without_prefix = module.get_name() == "line_break"
-    }
+    // Creates a root module and prints it.
+    let mut root_module = Module::new("Starship Root", "The root module", None);
+    root_module.set_segments(
+        formatter
+            .parse(None)
+            .expect("Unexpected error returned in root format variables"),
+    );
+
+    let module_strings = root_module.ansi_strings_for_shell(context.shell);
+    write!(buf, "{}", ANSIStrings(&module_strings)).unwrap();
 
     buf
 }
@@ -72,16 +95,15 @@ pub fn explain(args: ArgMatches) {
         desc: String,
     }
 
-    let dont_print = vec!["line_break", "character"];
+    let dont_print = vec!["line_break"];
 
     let modules = compute_modules(&context)
         .into_iter()
         .filter(|module| !dont_print.contains(&module.get_name().as_str()))
         .map(|module| {
-            let ansi_strings = module.ansi_strings();
             let value = module.get_segments().join("");
             ModuleInfo {
-                value: ansi_term::ANSIStrings(&ansi_strings[1..ansi_strings.len() - 1]).to_string(),
+                value: ansi_term::ANSIStrings(&module.ansi_strings()).to_string(),
                 value_len: value.chars().count() + count_wide_chars(&value),
                 desc: module.get_description().to_owned(),
             }
@@ -132,11 +154,38 @@ pub fn explain(args: ArgMatches) {
 }
 
 fn compute_modules<'a>(context: &'a Context) -> Vec<Module<'a>> {
-    enum Mod<'a> {
-        Builtin(&'a str),
-        Custom(&'a str),
+    let mut prompt_order: Vec<Module<'a>> = Vec::new();
+
+    let config = context.config.get_root_config();
+    let formatter = if let Ok(formatter) = StringFormatter::new(config.format) {
+        formatter
+    } else {
+        log::error!("Error parsing `format`");
+        return Vec::new();
+    };
+    let modules = formatter.get_variables();
+
+    for module in &modules {
+        // Manually add all modules if `$all` is encountered
+        if module == "all" {
+            for module in PROMPT_ORDER.iter() {
+                let modules = handle_module(module, &context, &modules);
+                prompt_order.extend(modules.into_iter());
+            }
+        } else {
+            let modules = handle_module(module, &context, &modules);
+            prompt_order.extend(modules.into_iter());
+        }
     }
 
+    prompt_order
+}
+
+fn handle_module<'a>(
+    module: &str,
+    context: &'a Context,
+    module_list: &BTreeSet<String>,
+) -> Vec<Module<'a>> {
     struct DebugCustomModules<'tmp>(&'tmp toml::value::Table);
 
     impl Debug for DebugCustomModules<'_> {
@@ -145,74 +194,63 @@ fn compute_modules<'a>(context: &'a Context) -> Vec<Module<'a>> {
         }
     }
 
-    let mut prompt_order: Vec<Mod> = Vec::new();
+    let mut modules: Vec<Option<Module>> = Vec::new();
 
-    // Write out a custom prompt order
-    let config_prompt_order = context.config.get_root_config().prompt_order;
-
-    for module in &config_prompt_order {
-        if ALL_MODULES.contains(module) {
-            // Write out a module if it isn't disabled
-            if !context.is_module_disabled_in_config(*module) {
-                prompt_order.push(Mod::Builtin(module));
-            }
-        } else if *module == "custom" {
-            // Write out all custom modules, except for those that are explicitly set
-            if let Some(custom_modules) = context.config.get_custom_modules() {
-                for (custom_module, config) in custom_modules {
-                    if should_add_implicit_custom_module(
-                        custom_module,
-                        config,
-                        &config_prompt_order,
-                    ) {
-                        prompt_order.push(Mod::Custom(custom_module));
+    if ALL_MODULES.contains(&module) {
+        // Write out a module if it isn't disabled
+        if !context.is_module_disabled_in_config(module) {
+            modules.push(modules::handle(module, &context));
+        }
+    } else if module == "custom" {
+        // Write out all custom modules, except for those that are explicitly set
+        if let Some(custom_modules) = context.config.get_custom_modules() {
+            let custom_modules = custom_modules
+                .iter()
+                .map(|(custom_module, config)| {
+                    if should_add_implicit_custom_module(custom_module, config, &module_list) {
+                        modules::custom::module(custom_module, &context)
+                    } else {
+                        None
                     }
-                }
-            }
-        } else if module.starts_with("custom.") {
-            // Write out a custom module if it isn't disabled (and it exists...)
-            match context.is_custom_module_disabled_in_config(&module[7..]) {
-                Some(true) => (), // Module is disabled, we don't add it to the prompt
-                Some(false) => prompt_order.push(Mod::Custom(&module[7..])),
-                None => match context.config.get_custom_modules() {
-                    Some(modules) => log::debug!(
-                        "prompt_order contains custom module \"{}\", but no configuration was provided. Configuration for the following modules were provided: {:?}",
-                        module,
-                        DebugCustomModules(modules),
+                })
+                .collect::<Vec<Option<Module<'a>>>>();
+            modules.extend(custom_modules)
+        }
+    } else if module.starts_with("custom.") {
+        // Write out a custom module if it isn't disabled (and it exists...)
+        match context.is_custom_module_disabled_in_config(&module[7..]) {
+            Some(true) => (), // Module is disabled, we don't add it to the prompt
+            Some(false) => modules.push(modules::custom::module(&module[7..], &context)),
+            None => match context.config.get_custom_modules() {
+                Some(modules) => log::debug!(
+                    "prompt_order contains custom module \"{}\", but no configuration was provided. Configuration for the following modules were provided: {:?}",
+                    module,
+                    DebugCustomModules(modules),
                     ),
-                    None => log::debug!(
-                        "prompt_order contains custom module \"{}\", but no configuration was provided.",
-                        module,
+                None => log::debug!(
+                    "prompt_order contains custom module \"{}\", but no configuration was provided.",
+                    module,
                     ),
-                },
-            }
-        } else {
-            log::debug!(
-                "Expected prompt_order to contain value from {:?}. Instead received {}",
-                ALL_MODULES,
-                module,
-            );
+            },
         }
+    } else {
+        log::debug!(
+            "Expected prompt_order to contain value from {:?}. Instead received {}",
+            ALL_MODULES,
+            module,
+        );
     }
 
-    prompt_order
-        .par_iter()
-        .map(|module| match module {
-            Mod::Builtin(builtin) => modules::handle(builtin, context),
-            Mod::Custom(custom) => modules::custom::module(custom, context),
-        }) // Compute segments
-        .flatten() // Remove segments set to `None`
-        .collect::<Vec<Module<'a>>>()
+    modules.into_iter().flatten().collect()
 }
 
 fn should_add_implicit_custom_module(
     custom_module: &str,
     config: &toml::Value,
-    config_prompt_order: &[&str],
+    module_list: &BTreeSet<String>,
 ) -> bool {
-    let is_explicitly_specified = config_prompt_order.iter().any(|x| {
-        x.len() == 7 + custom_module.len() && &x[..7] == "custom." && &x[7..] == custom_module
-    });
+    let explicit_module_name = format!("custom.{}", custom_module);
+    let is_explicitly_specified = module_list.contains(&explicit_module_name);
 
     if is_explicitly_specified {
         // The module is already specified explicitly, so we skip it