48094: apply cleanup commands of the current theme on any theme change

* 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.

3 files changed, 42 insertions, 19 deletions

diff --git a/ChangeLog b/ChangeLog
index cee131eda..7e96cfadf 100644
--- a/ChangeLog
+++ b/ChangeLog
@@ -1,3 +1,9 @@
+2021-02-22  Bart Schaefer  <schaefer@zsh.org>
+
+	* 48094: Arseny Maslennikov: Functions/Prompts/prompt_restore_setup,
+	Functions/Prompts/promptinit: apply cleanup commands of the current
+	theme on any theme change, including "prompt -[hp] $theme"
+
 2021-02-19  dana  <dana@dana.is>
 
 	* 47867: Arseny Maslennikov: Functions/Prompts/promptinit: Fix
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[@]}"