feat(config): Adds support for --profile <custom profile name> (#3467)

Co-authored-by: Kevin Song <chips@ksong.dev>
Co-authored-by: David Knaack <davidkna@users.noreply.github.com>

5 files changed, 115 insertions, 20 deletions

diff --git a/.github/config-schema.json b/.github/config-schema.json
index 702d57cb8..b23233787 100644
--- a/.github/config-schema.json
+++ b/.github/config-schema.json
@@ -1744,6 +1744,13 @@
           "type": "string"
         }
       }
+    },
+    "profiles": {
+      "default": {},
+      "type": "object",
+      "additionalProperties": {
+        "type": "string"
+      }
     }
   },
   "additionalProperties": false,
diff --git a/src/configs/starship_root.rs b/src/configs/starship_root.rs
index d857ae55f..3f29d0f6b 100644
--- a/src/configs/starship_root.rs
+++ b/src/configs/starship_root.rs
@@ -1,3 +1,4 @@
+use indexmap::IndexMap;
 use serde::{Deserialize, Serialize};
 use std::collections::HashMap;
 
@@ -20,6 +21,7 @@ pub struct StarshipRootConfig {
     #[serde(skip_serializing_if = "Option::is_none")]
     pub palette: Option<String>,
     pub palettes: HashMap<String, Palette>,
+    pub profiles: IndexMap<String, String>,
 }
 
 pub type Palette = HashMap<String, String>;
@@ -125,6 +127,7 @@ impl Default for StarshipRootConfig {
             format: "$all".to_string(),
             right_format: String::new(),
             continuation_prompt: "[∙](bright-black) ".to_string(),
+            profiles: Default::default(),
             scan_timeout: 30,
             command_timeout: 500,
             add_newline: true,
diff --git a/src/context.rs b/src/context.rs
index 702837647..0a20ea7c6 100644
--- a/src/context.rs
+++ b/src/context.rs
@@ -647,11 +647,12 @@ pub enum Shell {
 }
 
 /// Which kind of prompt target to print (main prompt, rprompt, ...)
-#[derive(Debug, Clone, Copy, PartialEq, Eq)]
+#[derive(Debug, Clone, PartialEq, Eq)]
 pub enum Target {
     Main,
     Right,
     Continuation,
+    Profile(String),
 }
 
 /// Properties as passed on from the shell as arguments
diff --git a/src/main.rs b/src/main.rs
index 92cbe8c17..17312f86c 100644
--- a/src/main.rs
+++ b/src/main.rs
@@ -85,8 +85,11 @@ enum Commands {
         /// Print the right prompt (instead of the standard left prompt)
         #[clap(long)]
         right: bool,
-        /// Print the continuation prompt (instead of the standard left prompt)
+        /// Print the prompt with the specified profile name (instead of the standard left prompt)
         #[clap(long, conflicts_with = "right")]
+        profile: Option<String>,
+        /// Print the continuation prompt (instead of the standard left prompt)
+        #[clap(long, conflicts_with = "right", conflicts_with = "profile")]
         continuation: bool,
         #[clap(flatten)]
         properties: Properties,
@@ -172,12 +175,14 @@ fn main() {
         Commands::Prompt {
             properties,
             right,
+            profile,
             continuation,
         } => {
-            let target = match (right, continuation) {
-                (true, _) => Target::Right,
-                (_, true) => Target::Continuation,
-                (_, _) => Target::Main,
+            let target = match (right, profile, continuation) {
+                (true, _, _) => Target::Right,
+                (_, Some(profile_name), _) => Target::Profile(profile_name),
+                (_, _, true) => Target::Continuation,
+                (_, _, _) => Target::Main,
             };
             print::prompt(properties, target)
         }
diff --git a/src/print.rs b/src/print.rs
index 5702cf4e0..6ff16d61d 100644
--- a/src/print.rs
+++ b/src/print.rs
@@ -11,6 +11,7 @@ use unicode_width::UnicodeWidthChar;
 
 use crate::configs::PROMPT_ORDER;
 use crate::context::{Context, Properties, Shell, Target};
+use crate::formatter::string_formatter::StringFormatterError;
 use crate::formatter::{StringFormatter, VariableHolder};
 use crate::module::Module;
 use crate::module::ALL_MODULES;
@@ -403,32 +404,42 @@ fn all_modules_uniq(module_list: &BTreeSet<String>) -> Vec<String> {
 /// and the list of all modules used in a format string
 fn load_formatter_and_modules<'a>(context: &'a Context) -> (StringFormatter<'a>, BTreeSet<String>) {
     let config = &context.root_config;
+    let (formatter, config_param) = match &context.target {
+        Target::Main => (StringFormatter::new(&config.format), "format".to_string()),
+        Target::Right => (
+            StringFormatter::new(&config.right_format),
+            "right_format".to_string(),
+        ),
+        Target::Continuation => (
+            StringFormatter::new(&config.continuation_prompt),
+            "continuation_prompt".to_string(),
+        ),
+        Target::Profile(name) => (
+            match config.profiles.get(name) {
+                Some(format) => StringFormatter::new(format),
+                _ => Err(StringFormatterError::Custom("Invalid Profile".to_string())),
+            },
+            format!("profile: {}", &name),
+        ),
+    };
 
-    let lformatter = StringFormatter::new(&config.format);
     let rformatter = StringFormatter::new(&config.right_format);
-    let cformatter = StringFormatter::new(&config.continuation_prompt);
-    if lformatter.is_err() {
-        log::error!("Error parsing `format`")
+
+    if formatter.is_err() {
+        log::error!("Error parsing `{}`", config_param);
     }
     if rformatter.is_err() {
         log::error!("Error parsing `right_format`")
     }
-    if cformatter.is_err() {
-        log::error!("Error parsing `continuation_prompt`")
-    }
 
-    match (lformatter, rformatter, cformatter) {
-        (Ok(lf), Ok(rf), Ok(cf)) => {
+    match (formatter, rformatter) {
+        (Ok(lf), Ok(rf)) => {
             let mut modules: BTreeSet<String> = BTreeSet::new();
             if context.target != Target::Continuation {
                 modules.extend(lf.get_variables());
                 modules.extend(rf.get_variables());
             }
-            match context.target {
-                Target::Main => (lf, modules),
-                Target::Right => (rf, modules),
-                Target::Continuation => (cf, modules),
-            }
+            (lf, modules)
         }
         _ => (StringFormatter::raw(">"), BTreeSet::new()),
     }
@@ -476,6 +487,26 @@ mod test {
     use crate::test::default_context;
 
     #[test]
+    fn main_prompt() {
+        let mut context = default_context();
+        context.config = StarshipConfig {
+            config: Some(toml::toml! {
+                add_newline=false
+                format="$character"
+                [character]
+                format=">\n>"
+            }),
+        };
+        context.root_config.format = "$character".to_string();
+        context.target = Target::Main;
+        context.root_config.add_newline = false;
+
+        let expected = String::from(">\n>");
+        let actual = get_prompt(context);
+        assert_eq!(expected, actual);
+    }
+
+    #[test]
     fn right_prompt() {
         let mut context = default_context();
         context.config = StarshipConfig {
@@ -494,6 +525,54 @@ mod test {
     }
 
     #[test]
+    fn custom_prompt() {
+        let mut context = default_context();
+        context.config = StarshipConfig {
+            config: Some(toml::toml! {
+                add_newline = false
+                [profiles]
+                test="0_0$character"
+                [character]
+                format=">>"
+            }),
+        };
+        context
+            .root_config
+            .profiles
+            .insert("test".to_string(), "0_0$character".to_string());
+        context.target = Target::Profile("test".to_string());
+        context.root_config.add_newline = false;
+
+        let expected = String::from("0_0>>");
+        let actual = get_prompt(context);
+        assert_eq!(expected, actual);
+    }
+
+    #[test]
+    fn custom_prompt_fallback() {
+        let mut context = default_context();
+        context.config = StarshipConfig {
+            config: Some(toml::toml! {
+                add_newline=false
+                [profiles]
+                test="0_0$character"
+                [character]
+                format=">>"
+            }),
+        };
+        context
+            .root_config
+            .profiles
+            .insert("test".to_string(), "0_0$character".to_string());
+        context.target = Target::Profile("wrong_prompt".to_string());
+        context.root_config.add_newline = false;
+
+        let expected = String::from(">");
+        let actual = get_prompt(context);
+        assert_eq!(expected, actual);
+    }
+
+    #[test]
     fn continuation_prompt() {
         let mut context = default_context();
         context.config = StarshipConfig {