feat: Display Vi mode as PROMPT (#169)

Add Vi-mode indicator for zsh

4 files changed, 21 insertions, 13 deletions

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")
-    }
-
 }