[PATCH v2] promptinit: Fix prompt cleanups

[Arseny Maslennikov <ar@xxxxxxxxx>]

The Zsh documentation at Doc/Zsh/contrib, paragraph 26.6.4 "Writing
Themes", says:

Declare cleanup

    If your function makes any other changes that should be undone when
    the theme is disabled, your setup function may call

    prompt_cleanup command

    where command should be suitably quoted. If your theme is ever
    disabled or replaced by another, command is executed with eval. You
    may declare more than one such cleanup hook.

The promptinit framework fails to apply cleanup commands of the current
theme on any theme change other than "prompt restore", as well as when
invoking setup functions of a certain $theme as an implementation detail
of `prompt -[hp] $theme'.

This problem was previously brought up in zsh-users/23314, to no avail.

We fix it in the following way, hopefully without breaking compatibility:
* Rename zstyle `cleanup' on the context `:prompt-theme' to `restore'
  everywhere but in prompt_cleanup(). It is only used as a restore
  mechanism now.
* Ensure prompt_cleanup() continues to store its command list in the
  `cleanup' style.
* Clean up before theme switch at the end of set_prompt().
* Prepend every use of prompt_*_setup (which might modify the shell
  state in ways that require cleanup) with a cleanup run.
* Adjust `prompt restore' to do both parts of the newly split restore
  mechanism, cleanup first.
---

Changes since v1:
* fix the case when no theme was previously activated
* small refactorings
---

 Functions/Prompts/prompt_restore_setup |  1 +
 Functions/Prompts/promptinit           | 54 +++++++++++++++++---------
 2 files changed, 36 insertions(+), 19 deletions(-)

diff --git a/Functions/Prompts/prompt_restore_setup b/Functions/Prompts/prompt_restore_setup
index 54c4adbf9..b77dbe815 100644
--- a/Functions/Prompts/prompt_restore_setup
+++ b/Functions/Prompts/prompt_restore_setup
@@ -1,2 +1,3 @@
 # Damn that was easy
 zstyle -t :prompt-theme cleanup
+zstyle -t :prompt-theme restore
diff --git a/Functions/Prompts/promptinit b/Functions/Prompts/promptinit
index 5e42ebdd3..37d69f100 100644
--- a/Functions/Prompts/promptinit
+++ b/Functions/Prompts/promptinit
@@ -47,17 +47,19 @@ prompt_preview_safely() {
     return
   fi
 
-  # This handles all the stuff from the default :prompt-theme cleanup
+  # This handles all the stuff from the default :prompt-theme restore
   local +h PS1=$PS1 PS2=$PS2 PS3=$PS3 PS4=$PS4 RPS1=$RPS1 RPS2=$RPS2
   local +h PROMPT=$PROMPT RPROMPT=$RPROMPT RPROMPT2=$RPROMPT2 PSVAR=$PSVAR
-  local -a precmd_functions preexec_functions prompt_preview_cleanup
+  local -a precmd_functions preexec_functions prompt_preview_restore
   local -aLl +h zle_highlight
 
   {
     # Save and clear current restore-point if any
-    zstyle -g prompt_preview_cleanup :prompt-theme cleanup
+    zstyle -g prompt_preview_restore :prompt-theme restore
     {
-      zstyle -d :prompt-theme cleanup
+      zstyle -d :prompt-theme restore
+      # Execute current cleanup sequence, if any.
+      zstyle -t :prompt-theme cleanup
 
       # The next line is a bit ugly.  It (perhaps unnecessarily)
       # runs the prompt theme setup function to ensure that if
@@ -74,8 +76,8 @@ prompt_preview_safely() {
       zstyle -t :prompt-theme cleanup
     }
   } always {
-    (( $#prompt_preview_cleanup )) &&
-      zstyle -e :prompt-theme cleanup "${prompt_preview_cleanup[@]}"
+    (( $#prompt_preview_restore )) &&
+      zstyle -e :prompt-theme restore "${prompt_preview_restore[@]}"
   }
 }
 
@@ -103,9 +105,11 @@ Use prompt -h <theme> for help on specific themes.'
         local +h PS1=$PS1 PS2=$PS2 PS3=$PS3 PS4=$PS4 RPS1=$RPS1 RPS2=$RPS2
         local +h PROMPT=$PROMPT RPROMPT=$RPROMPT RPROMPT2=$RPROMPT2 PSVAR=$PSVAR
         local -a precmd_functions preexec_functions
+        local theme_reset=''
       else
-        trap 'prompt_${prompt_theme[1]}_setup "${(@)prompt_theme[2,-1]}"' 0
+        local theme_reset='prompt_${prompt_theme[1]}_setup "${(@)prompt_theme[2,-1]}"'
       fi
+      trap 'zstyle -t :prompt-theme cleanup;'"${theme_reset:+ $theme_reset}" 0
       ;;
   esac
   case "$opt" in
@@ -120,6 +124,7 @@ Use prompt -h <theme> for help on specific themes.'
        ;;
     h) if [[ -n "$2" && -n $prompt_themes[(r)$2] ]]; then
          if functions prompt_$2_setup >/dev/null; then
+           zstyle -t :prompt-theme cleanup
            # The next line is a bit ugly.  It (perhaps unnecessarily)
            # runs the prompt theme setup function to ensure that if
            # the theme has a _help function that it's been autoloaded.
@@ -179,28 +184,38 @@ Use prompt -h <theme> for help on specific themes.'
        typeset -ga zle_highlight=( ${zle_highlight:#default:*} )
        (( ${#zle_highlight} )) || unset zle_highlight
 
+       zstyle -t :prompt-theme cleanup
        prompt_$1_setup "$@[2,-1]" && prompt_theme=( "$@" )
        ;;
   esac
 }
 
 prompt_cleanup () {
-  local -a cleanup_hooks
-  if zstyle -g cleanup_hooks :prompt-theme cleanup
-  then
-    cleanup_hooks+=(';' "$@")
-    zstyle -e :prompt-theme cleanup "${cleanup_hooks[@]}"
-  elif (( $+prompt_preview_cleanup == 0 ))
+  local -a cleanup_hooks theme_active
+  if ! zstyle -g cleanup_hooks :prompt-theme cleanup
   then
-    print -u2 "prompt_cleanup: no prompt theme active"
-    return 1
+    (( $+prompt_preview_restore == 0 )) &&
+    if ! zstyle -g theme_active :prompt-theme restore
+    then
+      print -u2 "prompt_cleanup: no prompt theme active"
+      return 1
+    fi
+
+    # Set the cleanup sequence up.
+    zstyle -e :prompt-theme cleanup \
+        'zstyle -d :prompt-theme cleanup;' \
+        'reply=(yes)'
+    zstyle -g cleanup_hooks :prompt-theme cleanup
   fi
+
+  cleanup_hooks+=(';' "$@")
+  zstyle -e :prompt-theme cleanup "${cleanup_hooks[@]}"
 }
 
 prompt () {
   local -a prompt_opts theme_active
 
-  zstyle -g theme_active :prompt-theme cleanup || {
+  zstyle -g theme_active :prompt-theme restore || {
     # This is done here rather than in set_prompt so that it
     # is safe and sane for set_prompt to setopt localoptions,
     # which will be cleared before we arrive back here again.
@@ -210,8 +225,8 @@ prompt () {
     [[ -o promptpercent ]] && prompt_opts+=(percent)
     [[ -o promptsp ]] && prompt_opts+=(sp)
     [[ -o promptsubst ]] && prompt_opts+=(subst)
-    zstyle -e :prompt-theme cleanup \
-        'zstyle -d :prompt-theme cleanup;' \
+    zstyle -e :prompt-theme restore \
+        'zstyle -d :prompt-theme restore;' \
 	'prompt_default_setup;' \
         ${PS1+PS1="${(q)PS1}"} \
         ${PS2+PS2="${(q)PS2}"} \
@@ -239,7 +254,7 @@ prompt_preview_theme () {
   emulate -L zsh
 
   # Check for proper state handling
-  (( $+prompt_preview_cleanup )) || {
+  (( $+prompt_preview_restore )) || {
     prompt_preview_safely "$@"
     return
   }
@@ -249,6 +264,7 @@ prompt_preview_theme () {
   print -n "$1 theme"
   (( $#* > 1 )) && print -n " with parameters \`$*[2,-1]'"
   print ":"
+  zstyle -t :prompt-theme cleanup
   prompt_${1}_setup "$@[2,-1]"
   (( ${#prompt_opts} )) &&
       setopt noprompt{bang,cr,percent,sp,subst} "prompt${^prompt_opts[@]}"
-- 
2.30.1