From 1a72757f019f59362c34ec81bf2811eb3d1d70da Mon Sep 17 00:00:00 2001 From: Daniel Watkins Date: Sat, 6 Apr 2024 09:28:26 -0400 Subject: [PATCH] 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). --- src/module.rs | 24 +++--------------------- src/print.rs | 9 +++++++-- src/utils.rs | 27 +++++++++------------------ 3 files changed, 19 insertions(+), 41 deletions(-) diff --git a/src/module.rs b/src/module.rs index abcb25f7..c067aa3b 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 { - 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) -> Vec { + pub fn ansi_strings_for_width(&self, width: Option) -> Vec { let mut iter = self.segments.iter().peekable(); let mut ansi_strings: Vec = 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, shell: Shell) -> Vec { - ansi_strings - .into_iter() - .map(|ansi| { - let wrapped = wrap_colorseq_for_shell(ansi.to_string(), shell); - AnsiString::from(wrapped) - }) - .collect::>() -} - fn ansi_line<'a, I>(segments: &mut I, term_width: Option) -> Vec> where I: Iterator, diff --git a/src/print.rs b/src/print.rs index 9ef18f32..a39a961f 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 eafeca54..c2c54b62 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() }