fix: Solve bash & zsh cursor location confusion issue (#183)

Solves the issue of cursor location confusion on bash and zsh (#110) . Solution: modify ANSIString and wrap non-printing characters with correct escape sequences.

2 files changed, 52 insertions, 1 deletions

diff --git a/src/init.rs b/src/init.rs
index 5d2ad65f8..1479d2174 100644
--- a/src/init.rs
+++ b/src/init.rs
@@ -92,6 +92,7 @@ starship_preexec() {
 };
 starship_precmd() {
     STATUS=$?;
+    export STARSHIP_SHELL="bash";
     "${starship_precmd_user_func-:}";
     if [[ $STARSHIP_START_TIME ]]; then
         STARSHIP_END_TIME=$(date +%s);
@@ -140,6 +141,7 @@ fi;
 const ZSH_INIT: &str = r##"
 starship_precmd() {
     STATUS=$?;
+    export STARSHIP_SHELL="zsh";
     if [[ $STARSHIP_START_TIME ]]; then
         STARSHIP_END_TIME="$(date +%s)";
         STARSHIP_DURATION=$((STARSHIP_END_TIME - STARSHIP_START_TIME));
diff --git a/src/module.rs b/src/module.rs
index 4c5787573..98f586cb3 100644
--- a/src/module.rs
+++ b/src/module.rs
@@ -79,12 +79,19 @@ 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> {
-        let mut ansi_strings = self
+        let shell = std::env::var("STARSHIP_SHELL").unwrap_or_default();
+        let ansi_strings = self
             .segments
             .iter()
             .map(Segment::ansi_string)
             .collect::<Vec<ANSIString>>();
 
+        let mut ansi_strings = match shell.as_str() {
+            "bash" => ansi_strings_modified(ansi_strings, shell),
+            "zsh" => ansi_strings_modified(ansi_strings, shell),
+            _ => ansi_strings,
+        };
+
         ansi_strings.insert(0, self.prefix.ansi_string());
         ansi_strings.push(self.suffix.ansi_string());
 
@@ -118,6 +125,48 @@ 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: String) -> Vec<ANSIString> {
+    const ESCAPE_BEGIN: char = '\u{1b}';
+    const MAYBE_ESCAPE_END: char = 'm';
+    ansi_strings
+        .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.as_str() {
+                            "bash" => String::from("\u{5c}\u{5b}\u{1b}"), // => \[ESC
+                            "zsh" => String::from("\u{25}\u{7b}\u{1b}"),  // => %{ESC
+                            _ => x.to_string(),
+                        }
+                    }
+                    MAYBE_ESCAPE_END => {
+                        if escaped {
+                            escaped = false;
+                            match shell.as_str() {
+                                "bash" => String::from("m\u{5c}\u{5d}"), // => m\]
+                                "zsh" => String::from("m\u{25}\u{7d}"),  // => m%}
+                                _ => x.to_string(),
+                            }
+                        } else {
+                            x.to_string()
+                        }
+                    }
+                    _ => x.to_string(),
+                })
+                .collect();
+            ANSIString::from(final_string)
+        })
+        .collect::<Vec<ANSIString>>()
+}
+
 /// Module affixes are to be used for the prefix or suffix of a module.
 pub struct Affix {
     /// The affix's name, to be used in configuration and logging.