fix: combine ANSI color codes before wrapping them (#5762)

* combine ANSI color codes before wrapping them

The existing code wraps each individual module's output for
`context.shell`, concatenates all that output together and passes it to
`AnsiStrings` to merge ANSI color codes.  However, the wrapping obscures
ANSI color codes, meaning that no merging is possible.

This commit changes the shell-specific wrapping to happen right before
output, once all modules' output has been concatenated together.  This
results in ANSI color codes being correctly merged, as well as reducing
the number of calls to `wrap_colorseq_for_shell` to one.

With a minimal `starship.toml`:

```
format = """$directory"""

[directory]
format = '[a]($style)[b]($style)'
```

The current code produces[0]:

```
\n%{\x1b[31m%}a%{\x1b[0m%}%{\x1b[31m%}b%{\x1b[0m%
```

And this commit's code:

```
\n%{\x1b[31m%}ab%{\x1b[0m%}
```

You can see that the current code emits an additional reset and repeated
color code between "a" and "b" compared to the new code.

[0] Produced in a Python shell with:

```
subprocess.check_output(
    "./target/debug/starship prompt", shell=True,
    env={"STARSHIP_CONFIG": "./starship.toml", "STARSHIP_SHELL": "zsh"}
)
```

* utils: return early from wrap_seq_for_shell unless wrapping required

* refactor(utils): simplify wrap_seq_for_shell

This commit modifies wrap_seq_for_shell to (a) return early for shells
with no wrapping required, and (b) determine the wrapping characters
once at the start of the function (rather than inline in the map
function for every character).

3 files changed, 19 insertions, 41 deletions

diff --git a/src/module.rs b/src/module.rs
index abcb25f79..c067aa3b2 100644
--- a/src/module.rs
+++ b/src/module.rs
@@ -1,7 +1,5 @@
-use crate::context::Shell;
 use crate::segment;
 use crate::segment::{FillSegment, Segment};
-use crate::utils::wrap_colorseq_for_shell;
 use nu_ansi_term::{AnsiString, AnsiStrings};
 use std::fmt;
 use std::time::Duration;
@@ -164,22 +162,16 @@ impl<'a> Module<'a> {
     /// Returns a vector of colored `AnsiString` elements to be later used with
     /// `AnsiStrings()` to optimize ANSI codes
     pub fn ansi_strings(&self) -> Vec<AnsiString> {
-        self.ansi_strings_for_shell(Shell::Unknown, None)
+        self.ansi_strings_for_width(None)
     }
 
-    pub fn ansi_strings_for_shell(&self, shell: Shell, width: Option<usize>) -> Vec<AnsiString> {
+    pub fn ansi_strings_for_width(&self, width: Option<usize>) -> Vec<AnsiString> {
         let mut iter = self.segments.iter().peekable();
         let mut ansi_strings: Vec<AnsiString> = Vec::new();
         while iter.peek().is_some() {
             ansi_strings.extend(ansi_line(&mut iter, width));
         }
-
-        match shell {
-            Shell::Bash => ansi_strings_modified(ansi_strings, shell),
-            Shell::Zsh => ansi_strings_modified(ansi_strings, shell),
-            Shell::Tcsh => ansi_strings_modified(ansi_strings, shell),
-            _ => ansi_strings,
-        }
+        ansi_strings
     }
 }
 
@@ -190,16 +182,6 @@ impl<'a> fmt::Display for Module<'a> {
     }
 }
 
-fn ansi_strings_modified(ansi_strings: Vec<AnsiString>, shell: Shell) -> Vec<AnsiString> {
-    ansi_strings
-        .into_iter()
-        .map(|ansi| {
-            let wrapped = wrap_colorseq_for_shell(ansi.to_string(), shell);
-            AnsiString::from(wrapped)
-        })
-        .collect::<Vec<AnsiString>>()
-}
-
 fn ansi_line<'a, I>(segments: &mut I, term_width: Option<usize>) -> Vec<AnsiString<'a>>
 where
     I: Iterator<Item = &'a Segment>,
diff --git a/src/print.rs b/src/print.rs
index 9ef18f321..a39a961f7 100644
--- a/src/print.rs
+++ b/src/print.rs
@@ -18,6 +18,7 @@ use crate::module::ALL_MODULES;
 use crate::modules;
 use crate::segment::Segment;
 use crate::shadow;
+use crate::utils::wrap_colorseq_for_shell;
 
 pub struct Grapheme<'a>(pub &'a str);
 
@@ -115,12 +116,16 @@ pub fn get_prompt(context: Context) -> String {
             .expect("Unexpected error returned in root format variables"),
     );
 
-    let module_strings = root_module.ansi_strings_for_shell(context.shell, Some(context.width));
+    let module_strings = root_module.ansi_strings_for_width(Some(context.width));
     if config.add_newline && context.target != Target::Continuation {
         // continuation prompts normally do not include newlines, but they can
         writeln!(buf).unwrap();
     }
-    write!(buf, "{}", AnsiStrings(&module_strings)).unwrap();
+    // AnsiStrings strips redundant ANSI color sequences, so apply it before modifying the ANSI
+    // color sequences for this specific shell
+    let shell_wrapped_output =
+        wrap_colorseq_for_shell(AnsiStrings(&module_strings).to_string(), context.shell);
+    write!(buf, "{}", shell_wrapped_output).unwrap();
 
     if context.target == Target::Right {
         // right prompts generally do not allow newlines
diff --git a/src/utils.rs b/src/utils.rs
index eafeca54b..c2c54b629 100644
--- a/src/utils.rs
+++ b/src/utils.rs
@@ -510,12 +510,13 @@ pub fn wrap_seq_for_shell(
     escape_begin: char,
     escape_end: char,
 ) -> String {
-    const BASH_BEG: &str = "\u{5c}\u{5b}"; // \[
-    const BASH_END: &str = "\u{5c}\u{5d}"; // \]
-    const ZSH_BEG: &str = "\u{25}\u{7b}"; // %{
-    const ZSH_END: &str = "\u{25}\u{7d}"; // %}
-    const TCSH_BEG: &str = "\u{25}\u{7b}"; // %{
-    const TCSH_END: &str = "\u{25}\u{7d}"; // %}
+    let (beg, end) = match shell {
+        // \[ and \]
+        Shell::Bash => ("\u{5c}\u{5b}", "\u{5c}\u{5d}"),
+        // %{ and %}
+        Shell::Tcsh | Shell::Zsh => ("\u{25}\u{7b}", "\u{25}\u{7d}"),
+        _ => return ansi,
+    };
 
     // ANSI escape codes cannot be nested, so we can keep track of whether we're
     // in an escape or not with a single boolean variable
@@ -525,20 +526,10 @@ pub fn wrap_seq_for_shell(
         .map(|x| {
             if x == escape_begin && !escaped {
                 escaped = true;
-                match shell {
-                    Shell::Bash => format!("{BASH_BEG}{escape_begin}"),
-                    Shell::Zsh => format!("{ZSH_BEG}{escape_begin}"),
-                    Shell::Tcsh => format!("{TCSH_BEG}{escape_begin}"),
-                    _ => x.to_string(),
-                }
+                format!("{beg}{escape_begin}")
             } else if x == escape_end && escaped {
                 escaped = false;
-                match shell {
-                    Shell::Bash => format!("{escape_end}{BASH_END}"),
-                    Shell::Zsh => format!("{escape_end}{ZSH_END}"),
-                    Shell::Tcsh => format!("{escape_end}{TCSH_END}"),
-                    _ => x.to_string(),
-                }
+                format!("{escape_end}{end}")
             } else {
                 x.to_string()
             }