diff --git a/src/context.rs b/src/context.rs index 14b7fa30..b8cea191 100644 --- a/src/context.rs +++ b/src/context.rs @@ -309,7 +309,7 @@ fn get_current_branch(repository: &Repository) -> Option { shorthand.map(std::string::ToString::to_string) } -#[derive(Debug, Clone, PartialEq)] +#[derive(Debug, Clone, Copy, PartialEq)] pub enum Shell { Bash, Fish, diff --git a/src/module.rs b/src/module.rs index 60108d47..81ecad51 100644 --- a/src/module.rs +++ b/src/module.rs @@ -1,6 +1,7 @@ use crate::config::SegmentConfig; use crate::context::Shell; use crate::segment::Segment; +use crate::utils::wrap_colorseq_for_shell; use ansi_term::Style; use ansi_term::{ANSIString, ANSIStrings}; use std::fmt; @@ -171,43 +172,12 @@ impl<'a> fmt::Display for Module<'a> { } } -/// Many shells cannot deal with raw unprintable characters (like ANSI escape sequences) and -/// miscompute the cursor position as a result, leading to strange visual bugs. Here, we wrap these -/// characters in shell-specific escape codes to indicate to the shell that they are zero-length. fn ansi_strings_modified(ansi_strings: Vec, shell: Shell) -> Vec { - const ESCAPE_BEGIN: char = '\u{1b}'; - const MAYBE_ESCAPE_END: char = 'm'; ansi_strings - .iter() + .into_iter() .map(|ansi| { - let mut escaped = false; - let final_string: String = ansi - .to_string() - .chars() - .map(|x| match x { - ESCAPE_BEGIN => { - escaped = true; - match shell { - Shell::Bash => String::from("\u{5c}\u{5b}\u{1b}"), // => \[ESC - Shell::Zsh => String::from("\u{25}\u{7b}\u{1b}"), // => %{ESC - _ => x.to_string(), - } - } - MAYBE_ESCAPE_END if escaped => { - escaped = false; - match shell { - Shell::Bash => String::from("m\u{5c}\u{5d}"), // => m\] - Shell::Zsh => String::from("m\u{25}\u{7d}"), // => m%} - _ => x.to_string(), - } - } - // escape the $ character to avoid $() code injection on bash shell, - // see #658 for more - '$' if Shell::Bash == shell => String::from("\\$"), - _ => x.to_string(), - }) - .collect(); - ANSIString::from(final_string) + let wrapped = wrap_colorseq_for_shell(ansi.to_string(), shell); + ANSIString::from(wrapped) }) .collect::>() } diff --git a/src/utils.rs b/src/utils.rs index d9bd2963..7cc26e42 100644 --- a/src/utils.rs +++ b/src/utils.rs @@ -3,6 +3,8 @@ use std::io::{Read, Result}; use std::path::Path; use std::process::Command; +use crate::context::Shell; + /// Return the string contents of a file pub fn read_file>(file_name: P) -> Result { let mut file = File::open(file_name)?; @@ -58,6 +60,55 @@ pub fn exec_cmd(cmd: &str, args: &[&str]) -> Option { } } +/// Wraps ANSI color escape sequences in the shell-appropriate wrappers. +pub fn wrap_colorseq_for_shell(ansi: String, shell: Shell) -> String { + const ESCAPE_BEGIN: char = '\u{1b}'; + const ESCAPE_END: char = 'm'; + wrap_seq_for_shell(ansi, shell, ESCAPE_BEGIN, ESCAPE_END) +} + +/// Many shells cannot deal with raw unprintable characters and miscompute the cursor position, +/// leading to strange visual bugs like duplicated/missing chars. This function wraps a specified +/// sequence in shell-specific escapes to avoid these problems. +pub fn wrap_seq_for_shell( + ansi: String, + shell: 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}"; // %} + + // 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 + let mut escaped = false; + let final_string: String = ansi + .chars() + .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), + _ => x.to_string(), + } + } else if x == escape_end && escaped { + escaped = false; + match shell { + Shell::Bash => format!("{}{}", escape_end, BASH_END), + Shell::Zsh => format!("{}{}", escape_end, ZSH_END), + _ => x.to_string(), + } + } else { + x.to_string() + } + }) + .collect(); + final_string +} + fn internal_exec_cmd(cmd: &str, args: &[&str]) -> Option { log::trace!("Executing command '{:?}' with args '{:?}'", cmd, args); match Command::new(cmd).args(args).output() { @@ -148,4 +199,42 @@ mod tests { assert_eq!(result, expected) } + + #[test] + fn test_color_sequence_wrappers() { + let test0 = "\x1b2mhellomynamekeyes\x1b2m"; // BEGIN: \x1b END: m + let test1 = "\x1b]330;mlol\x1b]0m"; // BEGIN: \x1b END: m + let test2 = "\u{1b}J"; // BEGIN: \x1b END: J + let test3 = "OH NO"; // BEGIN: O END: O + let test4 = "herpaderp"; + let test5 = ""; + + let zresult0 = wrap_seq_for_shell(test0.to_string(), Shell::Zsh, '\x1b', 'm'); + let zresult1 = wrap_seq_for_shell(test1.to_string(), Shell::Zsh, '\x1b', 'm'); + let zresult2 = wrap_seq_for_shell(test2.to_string(), Shell::Zsh, '\x1b', 'J'); + let zresult3 = wrap_seq_for_shell(test3.to_string(), Shell::Zsh, 'O', 'O'); + let zresult4 = wrap_seq_for_shell(test4.to_string(), Shell::Zsh, '\x1b', 'm'); + let zresult5 = wrap_seq_for_shell(test5.to_string(), Shell::Zsh, '\x1b', 'm'); + + assert_eq!(&zresult0, "%{\x1b2m%}hellomynamekeyes%{\x1b2m%}"); + assert_eq!(&zresult1, "%{\x1b]330;m%}lol%{\x1b]0m%}"); + assert_eq!(&zresult2, "%{\x1bJ%}"); + assert_eq!(&zresult3, "%{OH NO%}"); + assert_eq!(&zresult4, "herpaderp"); + assert_eq!(&zresult5, ""); + + let bresult0 = wrap_seq_for_shell(test0.to_string(), Shell::Bash, '\x1b', 'm'); + let bresult1 = wrap_seq_for_shell(test1.to_string(), Shell::Bash, '\x1b', 'm'); + let bresult2 = wrap_seq_for_shell(test2.to_string(), Shell::Bash, '\x1b', 'J'); + let bresult3 = wrap_seq_for_shell(test3.to_string(), Shell::Bash, 'O', 'O'); + let bresult4 = wrap_seq_for_shell(test4.to_string(), Shell::Bash, '\x1b', 'm'); + let bresult5 = wrap_seq_for_shell(test5.to_string(), Shell::Bash, '\x1b', 'm'); + + assert_eq!(&bresult0, "\\[\x1b2m\\]hellomynamekeyes\\[\x1b2m\\]"); + assert_eq!(&bresult1, "\\[\x1b]330;m\\]lol\\[\x1b]0m\\]"); + assert_eq!(&bresult2, "\\[\x1bJ\\]"); + assert_eq!(&bresult3, "\\[OH NO\\]"); + assert_eq!(&bresult4, "herpaderp"); + assert_eq!(&bresult5, ""); + } }