From 2134ecdc22ecba8041c5cce65ec64c190ad970c6 Mon Sep 17 00:00:00 2001
From: xenia <xenia@awoo.systems>
Date: Sun, 25 Sep 2022 16:57:31 -0400
Subject: [PATCH] monkeypatch for zsh-vi-mode

---
 src/context.rs           |  4 ++++
 src/init/starship.zsh    |  4 ++--
 src/modules/character.rs | 22 +++++++++++++++-------
 3 files changed, 21 insertions(+), 9 deletions(-)

diff --git a/src/context.rs b/src/context.rs
index a24a9889..a78fc330 100644
--- a/src/context.rs
+++ b/src/context.rs
@@ -863,6 +863,9 @@ pub struct Properties {
     /// The keymap of fish/zsh/cmd
     #[clap(short = 'k', long, default_value = "viins")]
     pub keymap: String,
+    // The ZVM_MODE variable in zsh-vi-mode
+    #[clap(short = 'z', long, default_value = "")]
+    pub zvm_mode: String,
     /// The number of currently running jobs
     #[clap(short, long, default_value_t, value_parser=parse_jobs)]
     pub jobs: i64,
@@ -878,6 +881,7 @@ impl Default for Properties {
             logical_path: None,
             cmd_duration: None,
             keymap: "viins".to_string(),
+            zvm_mode: "".to_string(),
             jobs: 0,
         }
     }
diff --git a/src/init/starship.zsh b/src/init/starship.zsh
index 2ca01052..5d19653a 100644
--- a/src/init/starship.zsh
+++ b/src/init/starship.zsh
@@ -87,6 +87,6 @@ VIRTUAL_ENV_DISABLE_PROMPT=1
 
 setopt promptsubst
 
-PROMPT='$('::STARSHIP::' prompt --terminal-width="$COLUMNS" --keymap="${KEYMAP:-}" --status="$STARSHIP_CMD_STATUS" --pipestatus="${STARSHIP_PIPE_STATUS[*]}" --cmd-duration="${STARSHIP_DURATION:-}" --jobs="$STARSHIP_JOBS_COUNT")'
-RPROMPT='$('::STARSHIP::' prompt --right --terminal-width="$COLUMNS" --keymap="${KEYMAP:-}" --status="$STARSHIP_CMD_STATUS" --pipestatus="${STARSHIP_PIPE_STATUS[*]}" --cmd-duration="${STARSHIP_DURATION:-}" --jobs="$STARSHIP_JOBS_COUNT")'
+PROMPT='$('::STARSHIP::' prompt --terminal-width="$COLUMNS" --keymap="${KEYMAP:-}" --zvm-mode="${ZVM_MODE:-}" --status="$STARSHIP_CMD_STATUS" --pipestatus="${STARSHIP_PIPE_STATUS[*]}" --cmd-duration="${STARSHIP_DURATION:-}" --jobs="$STARSHIP_JOBS_COUNT")'
+RPROMPT='$('::STARSHIP::' prompt --right --terminal-width="$COLUMNS" --keymap="${KEYMAP:-}" --zvm-mode="${ZVM_MODE:-}" --status="$STARSHIP_CMD_STATUS" --pipestatus="${STARSHIP_PIPE_STATUS[*]}" --cmd-duration="${STARSHIP_DURATION:-}" --jobs="$STARSHIP_JOBS_COUNT")'
 PROMPT2="$(::STARSHIP:: prompt --continuation)"
diff --git a/src/modules/character.rs b/src/modules/character.rs
index 232b3828..8cf67aa8 100644
--- a/src/modules/character.rs
+++ b/src/modules/character.rs
@@ -27,6 +27,7 @@ pub fn module<'a>(context: &'a Context) -> Option<Module<'a>> {
     let props = &context.properties;
     let exit_code = props.status_code.as_deref().unwrap_or("0");
     let keymap = props.keymap.as_str();
+    let zvm_mode = props.zvm_mode.as_str();
     let exit_success = exit_code == "0";
 
     // Match shell "keymap" names to normalized vi modes
@@ -34,13 +35,20 @@ pub fn module<'a>(context: &'a Context) -> Option<Module<'a>> {
     // Unfortunately, this is also the name of the non-vi default mode.
     // We do some environment detection in src/init.rs to translate.
     // The result: in non-vi fish, keymap is always reported as "insert"
-    let mode = match (&context.shell, keymap) {
-        (Shell::Fish, "default")
-        | (Shell::Zsh, "vicmd")
-        | (Shell::Cmd | Shell::PowerShell | Shell::Pwsh, "vi") => ShellEditMode::Normal,
-        (Shell::Fish, "visual") => ShellEditMode::Visual,
-        (Shell::Fish, "replace") => ShellEditMode::Replace,
-        (Shell::Fish, "replace_one") => ShellEditMode::ReplaceOne,
+    let mode = match (&context.shell, zvm_mode, keymap) {
+        (Shell::Zsh, "n", _) => ShellEditMode::Normal,
+        (Shell::Zsh, "i", _) => ShellEditMode::Insert,
+        (Shell::Zsh, "v", _) => ShellEditMode::Visual,
+        (Shell::Zsh, "vl", _) => ShellEditMode::Visual,
+        (Shell::Zsh, "r", _) => ShellEditMode::Replace,
+
+        (Shell::Fish, _, "default")
+            | (Shell::Zsh, _, "vicmd")
+            | (Shell::Cmd | Shell::PowerShell | Shell::Pwsh, _, "vi") => ShellEditMode::Normal,
+
+        (Shell::Fish, _, "visual") => ShellEditMode::Visual,
+        (Shell::Fish, _, "replace") => ShellEditMode::Replace,
+        (Shell::Fish, _, "replace_one") => ShellEditMode::ReplaceOne,
         _ => ASSUMED_MODE,
     };