From 190325049ef93731ab28295dbedf36d44ab33d7a Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Marc=20Cornell=C3=A0?= <hello@mcornella.com>
Date: Thu, 19 Aug 2021 18:20:57 +0200
Subject: [PATCH] fix(sudo): allow different $EDITOR settings and fix
 zsh-syntax-highlighting redraw

---
 plugins/sudo/sudo.plugin.zsh | 64 ++++++++++++++++++++++++++----------
 1 file changed, 46 insertions(+), 18 deletions(-)

diff --git a/plugins/sudo/sudo.plugin.zsh b/plugins/sudo/sudo.plugin.zsh
index f0a01bc1f..e02f88a87 100644
--- a/plugins/sudo/sudo.plugin.zsh
+++ b/plugins/sudo/sudo.plugin.zsh
@@ -25,6 +25,7 @@ __sudo-replace-buffer() {
 }
 
 sudo-command-line() {
+  # If line is empty, get the last run command from history
   [[ -z $BUFFER ]] && LBUFFER="$(fc -ln -1)"
 
   # Save beginning space
@@ -34,29 +35,56 @@ sudo-command-line() {
     LBUFFER="${LBUFFER:1}"
   fi
 
-  # Get the first part of the typed command and check if it's an alias to $EDITOR
-  # If so, locally change $EDITOR to the alias so that it matches below
-  if [[ -n "$EDITOR" ]]; then
-    local cmd="${${(Az)BUFFER}[1]}"
-    if [[ "${aliases[$cmd]} " = (\$EDITOR|$EDITOR)\ * ]]; then
-      local EDITOR="$cmd"
-    fi
-  fi
-
-  if [[ -n $EDITOR && $BUFFER = $EDITOR\ * ]]; then
-    __sudo-replace-buffer "$EDITOR" "sudoedit"
-  elif [[ -n $EDITOR && $BUFFER = \$EDITOR\ * ]]; then
-    __sudo-replace-buffer "\$EDITOR" "sudoedit"
-  elif [[ $BUFFER = sudoedit\ * ]]; then
-    __sudo-replace-buffer "sudoedit" "$EDITOR"
-  elif [[ $BUFFER = sudo\ * ]]; then
-    __sudo-replace-buffer "sudo" ""
+  # If $EDITOR is not set, just toggle the sudo prefix on and off
+  if [[ -z "$EDITOR" ]]; then
+    case "$BUFFER" in
+      sudoedit\ *) __sudo-replace-buffer "sudoedit" "" ;;
+      sudo\ *) __sudo-replace-buffer "sudo" "" ;;
+      *) LBUFFER="sudo $LBUFFER" ;;
+    esac
   else
-    LBUFFER="sudo $LBUFFER"
+    # Check if the typed command is really an alias to $EDITOR
+
+    # Get the first part of the typed command
+    local cmd="${${(Az)BUFFER}[1]}"
+    # Get the first part of the alias of the same name as $cmd, or $cmd if no alias matches
+    local realcmd="${${(Az)aliases[$cmd]}[1]:-$cmd}"
+    # Get the first part of the $EDITOR command ($EDITOR may have arguments after it)
+    local editorcmd="${${(Az)EDITOR}[1]}"
+
+    # Note: ${var:c} makes a $PATH search and expands $var to the full path
+    # The if condition is met when:
+    # - $realcmd is '$EDITOR'
+    # - $realcmd is "cmd" and $EDITOR is "cmd"
+    # - $realcmd is "cmd" and $EDITOR is "cmd --with --arguments"
+    # - $realcmd is "/path/to/cmd" and $EDITOR is "cmd"
+    # - $realcmd is "/path/to/cmd" and $EDITOR is "/path/to/cmd"
+    # or
+    # - $realcmd is "cmd" and $EDITOR is "cmd"
+    # - $realcmd is "cmd" and $EDITOR is "/path/to/cmd"
+    # or
+    # - $realcmd is "cmd" and $EDITOR is /alternative/path/to/cmd that appears in $PATH
+    if [[ "$realcmd" = (\$EDITOR|$editorcmd|${editorcmd:c}) \
+      || "${realcmd:c}" = ($editorcmd|${editorcmd:c}) ]] \
+      || builtin which -a "$realcmd" | command grep -Fx -q "$editorcmd"; then
+      editorcmd="$cmd" # replace $editorcmd with the typed command so it matches below
+    fi
+
+    # Check for editor commands in the typed command and replace accordingly
+    case "$BUFFER" in
+      $editorcmd\ *) __sudo-replace-buffer "$editorcmd" "sudoedit" ;;
+      \$EDITOR\ *) __sudo-replace-buffer '$EDITOR' "sudoedit" ;;
+      sudoedit\ *) __sudo-replace-buffer "sudoedit" "$EDITOR" ;;
+      sudo\ *) __sudo-replace-buffer "sudo" "" ;;
+      *) LBUFFER="sudo $LBUFFER" ;;
+    esac
   fi
 
   # Preserve beginning space
   LBUFFER="${WHITESPACE}${LBUFFER}"
+
+  # Redisplay edit buffer (compatibility with zsh-syntax-highlighting)
+  zle redisplay
 }
 
 zle -N sudo-command-line