refactor: Move ANSI escaping into its own functions (#866)

src/context.rs

@@ -309,7 +309,7 @@ fn get_current_branch(repository: &Repository) -> Option<String> {
     shorthand.map(std::string::ToString::to_string)
 }
 
-#[derive(Debug, Clone, PartialEq)]
+#[derive(Debug, Clone, Copy, PartialEq)]
 pub enum Shell {
     Bash,
     Fish,

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<ANSIString>, shell: Shell) -> Vec<ANSIString> {
-    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::<Vec<ANSIString>>()
 }

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<P: AsRef<Path>>(file_name: P) -> Result<String> {
     let mut file = File::open(file_name)?;
@@ -58,6 +60,55 @@ pub fn exec_cmd(cmd: &str, args: &[&str]) -> Option<CommandOutput> {
     }
 }
 
+/// 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<CommandOutput> {
     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, "");
+    }
 }