feat: Display Vi mode as PROMPT (#169)

Add Vi-mode indicator for zsh

6 files changed, 56 insertions, 13 deletions

diff --git a/docs/config/README.md b/docs/config/README.md
index ef327c9c7..8c3814f7d 100644
--- a/docs/config/README.md
+++ b/docs/config/README.md
@@ -100,6 +100,7 @@ can do this in two ways: by changing color (red/green) or by changing its shape
 | `symbol`                | `"➜"`   | The symbol used before the text input in the prompt. |
 | `error_symbol`          | `"✖"`   | The symbol used before text input if the previous command failed. |
 | `use_symbol_for_status` | `false` | Indicate error status by changing the symbol.         |
+| `vicmd_symbol`          | `"❮"`   | The symbol used before the text input in the prompt if zsh is in vim normal mode. |
 | `disabled`              | `false` | Disables the `character` module.                      |
 
 ### Example
diff --git a/src/init.rs b/src/init.rs
index d366bc1b8..4aed1f14d 100644
--- a/src/init.rs
+++ b/src/init.rs
@@ -158,6 +158,12 @@ if [[ ${preexec_functions[(ie)starship_preexec]} -gt ${#preexec_functions} ]]; t
     preexec_functions+=(starship_preexec);
 fi;
 STARSHIP_START_TIME="$(date +%s)";
+function zle-keymap-select
+{
+    PROMPT=$(starship prompt --keymap=$KEYMAP --jobs="$(jobs | wc -l)");
+    zle reset-prompt;
+};
+zle -N zle-keymap-select;
 "##;
 
 /* Fish setup is simple because they give us CMD_DURATION. Just account for name
diff --git a/src/main.rs b/src/main.rs
index c95b5fd21..523b2d9bc 100644
--- a/src/main.rs
+++ b/src/main.rs
@@ -43,6 +43,14 @@ fn main() {
         .help("The execution duration of the last command, in seconds")
         .takes_value(true);
 
+    let keymap_arg = Arg::with_name("keymap")
+        .short("k")
+        .long("keymap")
+        .value_name("KEYMAP")
+        // zsh only
+        .help("The keymap of zsh")
+        .takes_value(true);
+
     let jobs_arg = Arg::with_name("jobs")
         .short("j")
         .long("jobs")
@@ -69,6 +77,7 @@ fn main() {
                 .arg(&status_code_arg)
                 .arg(&path_arg)
                 .arg(&cmd_duration_arg)
+                .arg(&keymap_arg)
                 .arg(&jobs_arg),
         )
         .subcommand(
@@ -82,6 +91,7 @@ fn main() {
                 .arg(&status_code_arg)
                 .arg(&path_arg)
                 .arg(&cmd_duration_arg)
+                .arg(&keymap_arg)
                 .arg(&jobs_arg),
         )
         .get_matches();
diff --git a/src/modules/character.rs b/src/modules/character.rs
index 7946cff49..5e732dfa0 100644
--- a/src/modules/character.rs
+++ b/src/modules/character.rs
@@ -12,6 +12,8 @@ use ansi_term::Color;
 pub fn module<'a>(context: &'a Context) -> Option<Module<'a>> {
     const SUCCESS_CHAR: &str = "➜";
     const FAILURE_CHAR: &str = "✖";
+    const VICMD_CHAR: &str = "❮";
+
     let color_success = Color::Green.bold();
     let color_failure = Color::Red.bold();
 
@@ -23,11 +25,14 @@ pub fn module<'a>(context: &'a Context) -> Option<Module<'a>> {
         .config_value_bool("use_symbol_for_status")
         .unwrap_or(false);
     let exit_success = arguments.value_of("status_code").unwrap_or("0") == "0";
+    let keymap = arguments.value_of("keymap").unwrap_or("viins");
 
     /* If an error symbol is set in the config, use symbols to indicate
     success/failure, in addition to color */
     let symbol = if use_symbol && !exit_success {
         module.new_segment("error_symbol", FAILURE_CHAR)
+    } else if keymap == "vicmd" {
+        module.new_segment("vicmd_symbol", VICMD_CHAR)
     } else {
         module.new_segment("symbol", SUCCESS_CHAR)
     };
diff --git a/src/modules/python.rs b/src/modules/python.rs
index 8c9cc33a7..b651e1d51 100644
--- a/src/modules/python.rs
+++ b/src/modules/python.rs
@@ -111,17 +111,4 @@ mod tests {
         let input = "Python 3.7.2";
         assert_eq!(format_python_version(input), "v3.7.2");
     }
-
-    #[test]
-    fn test_no_virtual_env() {
-        env::set_var("VIRTUAL_ENV", "");
-        assert_eq!(get_python_virtual_env(), None)
-    }
-
-    #[test]
-    fn test_virtual_env() {
-        env::set_var("VIRTUAL_ENV", "/foo/bar/my_venv");
-        assert_eq!(get_python_virtual_env().unwrap(), "my_venv")
-    }
-
 }
diff --git a/tests/testsuite/character.rs b/tests/testsuite/character.rs
index 4c9fcaf5f..8913c9321 100644
--- a/tests/testsuite/character.rs
+++ b/tests/testsuite/character.rs
@@ -72,3 +72,37 @@ fn char_module_symbolyes_status() -> io::Result<()> {
 
     Ok(())
 }
+
+#[test]
+fn char_module_vicmd_keymap() -> io::Result<()> {
+    let expected_vicmd = format!("{} ", Color::Green.bold().paint("❮"));
+    let expected_specified = format!("{} ", Color::Green.bold().paint("N"));
+    let expected_other = format!("{} ", Color::Green.bold().paint("➜"));
+
+    // zle keymap is vicmd
+    let output = common::render_module("character")
+        .arg("--keymap=vicmd")
+        .output()?;
+    let actual = String::from_utf8(output.stdout).unwrap();
+    assert_eq!(expected_vicmd, actual);
+
+    // specified vicmd character
+    let output = common::render_module("character")
+        .use_config(toml::toml! {
+            [character]
+            vicmd_symbol = "N"
+        })
+        .arg("--keymap=vicmd")
+        .output()?;
+    let actual = String::from_utf8(output.stdout).unwrap();
+    assert_eq!(expected_specified, actual);
+
+    // zle keymap is other
+    let output = common::render_module("character")
+        .arg("--keymap=visual")
+        .output()?;
+    let actual = String::from_utf8(output.stdout).unwrap();
+    assert_eq!(expected_other, actual);
+
+    Ok(())
+}