From c4a126c2aca43d85ef6629ed5fc47137a5a530a4 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Marc=20Cornell=C3=A0?= <marc@mcornella.com>
Date: Thu, 30 May 2024 19:51:35 +0200
Subject: [PATCH] Initial async prompt utility wrapper

---
 lib/async_prompt.zsh          | 20 ++++++++++++++++++-
 lib/git.zsh                   | 37 +++++++++++++----------------------
 themes/minimal.zsh-theme      |  6 +++++-
 themes/robbyrussell.zsh-theme |  2 +-
 4 files changed, 39 insertions(+), 26 deletions(-)

diff --git a/lib/async_prompt.zsh b/lib/async_prompt.zsh
index db48446e7..6f1c5f4cc 100644
--- a/lib/async_prompt.zsh
+++ b/lib/async_prompt.zsh
@@ -26,7 +26,6 @@ autoload -Uz is-at-least
 # This API is subject to change and optimization. Rely on it at your own risk.
 
 function _omz_register_handler {
-  setopt localoptions noksharrays
   typeset -ga _omz_async_functions
   # we want to do nothing if there's no $1 function or we already set it up
   if [[ -z "$1" ]] || (( ! ${+functions[$1]} )) \
@@ -42,6 +41,25 @@ function _omz_register_handler {
   fi
 }
 
+function _omz_make_async_function {
+  local sync_func=$1
+
+  # Check if the sync function has already been made async
+  if (( ${+functions[${sync_func}_async]} )); then
+    return
+  fi
+
+  # Register the sync function as a handler
+  _omz_register_handler ${sync_func}
+
+  # Redefine the original function to output the async result
+  eval "function ${sync_func}_async {
+    if [[ -n \"\${_OMZ_ASYNC_OUTPUT[${sync_func}]}\" ]]; then
+      echo -n \"\${_OMZ_ASYNC_OUTPUT[${sync_func}]}\"
+    fi
+  }"
+}
+
 # Set up async handlers and callbacks
 function _omz_async_request {
   local -i ret=$?
diff --git a/lib/git.zsh b/lib/git.zsh
index b257d01a4..002c78552 100644
--- a/lib/git.zsh
+++ b/lib/git.zsh
@@ -11,7 +11,7 @@ function __git_prompt_git() {
   GIT_OPTIONAL_LOCKS=0 command git "$@"
 }
 
-function _omz_git_prompt_info() {
+function git_prompt_info() {
   # If we are on a folder not tracked by git, get out.
   # Otherwise, check for hide-info at global and local repository level
   if ! __git_prompt_git rev-parse --git-dir &> /dev/null \
@@ -41,31 +41,29 @@ function _omz_git_prompt_info() {
 
 # Use async version if setting is enabled or undefined
 if zstyle -T ':omz:alpha:lib:git' async-prompt; then
-  function git_prompt_info() {
-    if [[ -n "${_OMZ_ASYNC_OUTPUT[_omz_git_prompt_info]}" ]]; then
-      echo -n "${_OMZ_ASYNC_OUTPUT[_omz_git_prompt_info]}"
-    fi
-  }
-
-  function git_prompt_status() {
-    if [[ -n "${_OMZ_ASYNC_OUTPUT[_omz_git_prompt_status]}" ]]; then
-      echo -n "${_OMZ_ASYNC_OUTPUT[_omz_git_prompt_status]}"
-    fi
-  }
-
   # Conditionally register the async handler, only if it's needed in $PROMPT
   # or any of the other prompt variables
   function _defer_async_git_register() {
+    # Check if the user wants to register some async functions
+    local -a async_functions
+    zstyle -a ':omz:alpha:lib:git' async-functions async_functions
+
+    if (( ${#async_functions} )); then
+      for async_function in $async_functions; do
+        _omz_make_async_function $async_function
+      done
+    fi
+
     # Check if git_prompt_info is used in a prompt variable
     case "${PS1}:${PS2}:${PS3}:${PS4}:${RPROMPT}:${RPS1}:${RPS2}:${RPS3}:${RPS4}" in
     *(\$\(git_prompt_info\)|\`git_prompt_info\`)*)
-      _omz_register_handler _omz_git_prompt_info
+      _omz_make_async_function git_prompt_info
       ;;
     esac
 
     case "${PS1}:${PS2}:${PS3}:${PS4}:${RPROMPT}:${RPS1}:${RPS2}:${RPS3}:${RPS4}" in
     *(\$\(git_prompt_status\)|\`git_prompt_status\`)*)
-      _omz_register_handler _omz_git_prompt_status
+      _omz_make_async_function git_prompt_status
       ;;
     esac
 
@@ -76,13 +74,6 @@ if zstyle -T ':omz:alpha:lib:git' async-prompt; then
   # Register the async handler first. This needs to be done before
   # the async request prompt is run
   precmd_functions=(_defer_async_git_register $precmd_functions)
-else
-  function git_prompt_info() {
-    _omz_git_prompt_info
-  }
-  function git_prompt_status() {
-    _omz_git_prompt_status
-  }
 fi
 
 # Checks if working tree is dirty
@@ -213,7 +204,7 @@ function git_prompt_long_sha() {
   SHA=$(__git_prompt_git rev-parse HEAD 2> /dev/null) && echo "$ZSH_THEME_GIT_PROMPT_SHA_BEFORE$SHA$ZSH_THEME_GIT_PROMPT_SHA_AFTER"
 }
 
-function _omz_git_prompt_status() {
+function git_prompt_status() {
   [[ "$(__git_prompt_git config --get oh-my-zsh.hide-status 2>/dev/null)" = 1 ]] && return
 
   # Maps a git status prefix to an internal constant
diff --git a/themes/minimal.zsh-theme b/themes/minimal.zsh-theme
index 588ab6982..645dd0296 100644
--- a/themes/minimal.zsh-theme
+++ b/themes/minimal.zsh-theme
@@ -23,4 +23,8 @@ vcs_status() {
   fi
 }
 
-PROMPT='%2~ $(vcs_status)»%b '
+# Register the sync function
+autoload -Uz _omz_make_async_function
+_omz_make_async_function vcs_status
+
+PROMPT='%2~ $(vcs_status_async)»%b '
diff --git a/themes/robbyrussell.zsh-theme b/themes/robbyrussell.zsh-theme
index cfecfc892..ff6eb7c1a 100644
--- a/themes/robbyrussell.zsh-theme
+++ b/themes/robbyrussell.zsh-theme
@@ -1,5 +1,5 @@
 PROMPT="%(?:%{$fg_bold[green]%}%1{➜%} :%{$fg_bold[red]%}%1{➜%} ) %{$fg[cyan]%}%c%{$reset_color%}"
-PROMPT+=' $(git_prompt_info)'
+PROMPT+=' $(git_prompt_info_async)'
 
 ZSH_THEME_GIT_PROMPT_PREFIX="%{$fg_bold[blue]%}git:(%{$fg[red]%}"
 ZSH_THEME_GIT_PROMPT_SUFFIX="%{$reset_color%} "