From a3b92bb05360266c437f3bf793c948ecf63f778b Mon Sep 17 00:00:00 2001
From: Sorin Ionescu <sorin.ionescu@gmail.com>
Date: Wed, 21 Mar 2012 14:34:01 -0400
Subject: [PATCH] [Fix #51] Make the vi mode indicator more flexible

---
 keyboard.zsh                    | 43 ++++++++++-----------------------
 themes/sorin/prompt_sorin_setup |  5 ++--
 2 files changed, 16 insertions(+), 32 deletions(-)

diff --git a/keyboard.zsh b/keyboard.zsh
index 789f09d..38ce767 100644
--- a/keyboard.zsh
+++ b/keyboard.zsh
@@ -12,8 +12,11 @@ fi
 
 # The default styles.
 
+# Indicator to notify of vi insert mode.
+zstyle ':omz:prompt:vi' insert '>>>'
+
 # Indicator to notify of vi command mode.
-zstyle ':omz:prompt' vicmd '<<<'
+zstyle ':omz:prompt:vi' command '<<<'
 
 # Indicator to notify of generating completion.
 zstyle ':omz:completion' indicator '...'
@@ -61,39 +64,19 @@ keyinfo=(
   'BackTab'   "$terminfo[kcbt]"
 )
 
-# Restores RPROMPT when exiting vicmd.
-function vi-restore-rprompt() {
-  if (( $+RPROMPT_CACHED )); then
-    RPROMPT="$RPROMPT_CACHED"
-    unset RPROMPT_CACHED
-    zle reset-prompt
-    return 0
+# Displays the current vi mode.
+function zle-line-init zle-line-finish zle-keymap-select {
+  if [[ "$KEYMAP" == 'vicmd' ]]; then
+    zstyle -s ':omz:prompt:vi' command 'vi_prompt_info'
+  else
+    zstyle -s ':omz:prompt:vi' insert 'vi_prompt_info'
   fi
-  return 1
-}
-add-zsh-trap INT vi-restore-rprompt
-
-# Displays the current vi mode (command).
-function zle-keymap-select() {
-  if ! vi-restore-rprompt && [[ "$KEYMAP" == 'vicmd' ]]; then
-    RPROMPT_CACHED="$RPROMPT"
-    zstyle -s ':omz:prompt' vicmd RPROMPT
-    zle reset-prompt
-  fi
-}
-zle -N zle-keymap-select
-
-# Resets the prompt after exiting edit-command-line.
-function zle-line-init() {
-  vi-restore-rprompt
+  zle reset-prompt
+  zle -R
 }
 zle -N zle-line-init
-
-# Resets the prompt after the line has been accepted.
-function zle-line-finish() {
-  vi-restore-rprompt
-}
 zle -N zle-line-finish
+zle -N zle-keymap-select
 
 # Expands .... to ../..
 function expand-dot-to-parent-directory-path() {
diff --git a/themes/sorin/prompt_sorin_setup b/themes/sorin/prompt_sorin_setup
index 396d082..68c73e1 100644
--- a/themes/sorin/prompt_sorin_setup
+++ b/themes/sorin/prompt_sorin_setup
@@ -26,7 +26,8 @@ function prompt_sorin_setup() {
   add-zsh-hook precmd prompt_sorin_precmd
 
   zstyle ':omz:completion' indicator '%B%F{red}...%f%b'
-  zstyle ':omz:prompt' vicmd '%F{yellow}❮%f%B%F{red}❮%f%b%F{red}❮%f'
+  zstyle ':omz:prompt:vi' insert ''
+  zstyle ':omz:prompt:vi' command ' %F{yellow}❮%f%B%F{red}❮%f%b%F{red}❮%f'
   zstyle ':omz:plugin:git:prompt' action ':%%B%F{yellow}%s%f%%b'
   zstyle ':omz:plugin:git:prompt' added ' %%B%F{green}✚%f%%b'
   zstyle ':omz:plugin:git:prompt' ahead ' %%B%F{yellow}⬆%f%%b'
@@ -43,7 +44,7 @@ function prompt_sorin_setup() {
   zstyle ':omz:plugin:git:prompt' rprompt '%A%B%S%a%d%m%r%U%u'
 
   PROMPT='%F{cyan}%1~%f${(e)git_prompt_info} %(!.%B%F{red}#%f%b.%B%F{green}❯%f%b) '
-  RPROMPT='%(?::%F{red}⏎%f)${VIM:+" %B%F{green}V%f%b"}${git_rprompt_info}'
+  RPROMPT='${vi_prompt_info}%(?:: %F{red}⏎%f)${VIM:+" %B%F{green}V%f%b"}${git_rprompt_info}'
   SPROMPT='zsh: correct %F{red}%R%f to %F{green}%r%f [nyae]? '
 }