fix(vcs_info): avoid undefining wrapper during load transition

Co-authored-by: robbyrussell <257+robbyrussell@users.noreply.github.com>

.github/workflows/main.yml

@@ -36,8 +36,6 @@ jobs:
         run: |
           for file in ./oh-my-zsh.sh \
                       ./lib/*.zsh \
-                      ./functions/* \
-                      ./completions/_* \
                       ./plugins/*/*.plugin.zsh \
                       ./plugins/*/_* \
                       ./themes/*.zsh-theme; do

completions/_omz

@@ -1,90 +0,0 @@
-#compdef omz
-
-local -a cmds subcmds
-cmds=(
-  'changelog:Print the changelog'
-  'help:Usage information'
-  'plugin:Manage plugins'
-  'pr:Manage Oh My Zsh Pull Requests'
-  'reload:Reload the current zsh session'
-  'shop:Open the Oh My Zsh shop'
-  'theme:Manage themes'
-  'update:Update Oh My Zsh'
-  'version:Show the version'
-)
-
-if (( CURRENT == 2 )); then
-  _describe 'command' cmds
-elif (( CURRENT == 3 )); then
-  case "$words[2]" in
-    changelog) local -a refs
-      refs=("${(@f)$(builtin cd -q "$ZSH"; command git for-each-ref --format="%(refname:short):%(subject)" refs/heads refs/tags)}")
-      _describe 'command' refs ;;
-    plugin) subcmds=(
-      'disable:Disable plugin(s)'
-      'enable:Enable plugin(s)'
-      'info:Get plugin information'
-      'list:List plugins'
-      'load:Load plugin(s)'
-    )
-      _describe 'command' subcmds ;;
-    pr) subcmds=('clean:Delete all Pull Request branches' 'test:Test a Pull Request')
-      _describe 'command' subcmds ;;
-    theme) subcmds=('list:List themes' 'set:Set a theme in your .zshrc file' 'use:Load a theme')
-      _describe 'command' subcmds ;;
-  esac
-elif (( CURRENT == 4 )); then
-  case "${words[2]}::${words[3]}" in
-    plugin::(disable|enable|load))
-      local -aU valid_plugins
-
-      if [[ "${words[3]}" = disable ]]; then
-        # if command is "disable", only offer already enabled plugins
-        valid_plugins=($plugins)
-      else
-        valid_plugins=("$ZSH"/plugins/*/{_*,*.plugin.zsh}(-.N:h:t) "$ZSH_CUSTOM"/plugins/*/{_*,*.plugin.zsh}(-.N:h:t))
-        # if command is "enable", remove already enabled plugins
-        [[ "${words[3]}" = enable ]] && valid_plugins=(${valid_plugins:|plugins})
-      fi
-
-      _describe 'plugin' valid_plugins ;;
-    plugin::info)
-      local -aU plugins
-      plugins=("$ZSH"/plugins/*/{_*,*.plugin.zsh}(-.N:h:t) "$ZSH_CUSTOM"/plugins/*/{_*,*.plugin.zsh}(-.N:h:t))
-      _describe 'plugin' plugins ;;
-    plugin::list)
-      local -a opts
-      opts=('--enabled:List enabled plugins only')
-      _describe -o 'options' opts ;;
-    theme::(set|use))
-      local -aU themes
-      themes=("$ZSH"/themes/*.zsh-theme(-.N:t:r) "$ZSH_CUSTOM"/**/*.zsh-theme(-.N:r:gs:"$ZSH_CUSTOM"/themes/:::gs:"$ZSH_CUSTOM"/:::))
-      _describe 'theme' themes ;;
-  esac
-elif (( CURRENT > 4 )); then
-  case "${words[2]}::${words[3]}" in
-    plugin::(enable|disable|load))
-      local -aU valid_plugins
-
-      if [[ "${words[3]}" = disable ]]; then
-        # if command is "disable", only offer already enabled plugins
-        valid_plugins=($plugins)
-      else
-        valid_plugins=("$ZSH"/plugins/*/{_*,*.plugin.zsh}(-.N:h:t) "$ZSH_CUSTOM"/plugins/*/{_*,*.plugin.zsh}(-.N:h:t))
-        # if command is "enable", remove already enabled plugins
-        [[ "${words[3]}" = enable ]] && valid_plugins=(${valid_plugins:|plugins})
-      fi
-
-      # Remove plugins already passed as arguments
-      # NOTE: $(( CURRENT - 1 )) is the last plugin argument completely passed, i.e. that which
-      # has a space after them. This is to avoid removing plugins partially passed, which makes
-      # the completion not add a space after the completed plugin.
-      local -a args
-      args=(${words[4,$(( CURRENT - 1))]})
-      valid_plugins=(${valid_plugins:|args})
-
-      _describe 'plugin' valid_plugins ;;
-  esac
-fi
-
-return 0

lib/cli.zsh

@@ -1,3 +1,5 @@
+#!/usr/bin/env zsh
+
 function omz {
   setopt localoptions noksharrays
   [[ $# -gt 0 ]] || {
@@ -18,6 +20,102 @@ function omz {
   _omz::$command "$@"
 }
 
+function _omz {
+  local -a cmds subcmds
+  cmds=(
+    'changelog:Print the changelog'
+    'help:Usage information'
+    'plugin:Manage plugins'
+    'pr:Manage Oh My Zsh Pull Requests'
+    'reload:Reload the current zsh session'
+    'shop:Open the Oh My Zsh shop'
+    'theme:Manage themes'
+    'update:Update Oh My Zsh'
+    'version:Show the version'
+  )
+
+  if (( CURRENT == 2 )); then
+    _describe 'command' cmds
+  elif (( CURRENT == 3 )); then
+    case "$words[2]" in
+      changelog) local -a refs
+        refs=("${(@f)$(builtin cd -q "$ZSH"; command git for-each-ref --format="%(refname:short):%(subject)" refs/heads refs/tags)}")
+        _describe 'command' refs ;;
+      plugin) subcmds=(
+        'disable:Disable plugin(s)'
+        'enable:Enable plugin(s)'
+        'info:Get plugin information'
+        'list:List plugins'
+        'load:Load plugin(s)'
+      )
+        _describe 'command' subcmds ;;
+      pr) subcmds=('clean:Delete all Pull Request branches' 'test:Test a Pull Request')
+        _describe 'command' subcmds ;;
+      theme) subcmds=('list:List themes' 'set:Set a theme in your .zshrc file' 'use:Load a theme')
+        _describe 'command' subcmds ;;
+    esac
+  elif (( CURRENT == 4 )); then
+    case "${words[2]}::${words[3]}" in
+      plugin::(disable|enable|load))
+        local -aU valid_plugins
+
+        if [[ "${words[3]}" = disable ]]; then
+          # if command is "disable", only offer already enabled plugins
+          valid_plugins=($plugins)
+        else
+          valid_plugins=("$ZSH"/plugins/*/{_*,*.plugin.zsh}(-.N:h:t) "$ZSH_CUSTOM"/plugins/*/{_*,*.plugin.zsh}(-.N:h:t))
+          # if command is "enable", remove already enabled plugins
+          [[ "${words[3]}" = enable ]] && valid_plugins=(${valid_plugins:|plugins})
+        fi
+
+        _describe 'plugin' valid_plugins ;;
+      plugin::info)
+        local -aU plugins
+        plugins=("$ZSH"/plugins/*/{_*,*.plugin.zsh}(-.N:h:t) "$ZSH_CUSTOM"/plugins/*/{_*,*.plugin.zsh}(-.N:h:t))
+        _describe 'plugin' plugins ;;
+      plugin::list)
+        local -a opts
+        opts=('--enabled:List enabled plugins only')
+        _describe -o 'options' opts ;;
+      theme::(set|use))
+        local -aU themes
+        themes=("$ZSH"/themes/*.zsh-theme(-.N:t:r) "$ZSH_CUSTOM"/**/*.zsh-theme(-.N:r:gs:"$ZSH_CUSTOM"/themes/:::gs:"$ZSH_CUSTOM"/:::))
+        _describe 'theme' themes ;;
+    esac
+  elif (( CURRENT > 4 )); then
+    case "${words[2]}::${words[3]}" in
+      plugin::(enable|disable|load))
+        local -aU valid_plugins
+
+        if [[ "${words[3]}" = disable ]]; then
+          # if command is "disable", only offer already enabled plugins
+          valid_plugins=($plugins)
+        else
+          valid_plugins=("$ZSH"/plugins/*/{_*,*.plugin.zsh}(-.N:h:t) "$ZSH_CUSTOM"/plugins/*/{_*,*.plugin.zsh}(-.N:h:t))
+          # if command is "enable", remove already enabled plugins
+          [[ "${words[3]}" = enable ]] && valid_plugins=(${valid_plugins:|plugins})
+        fi
+
+        # Remove plugins already passed as arguments
+        # NOTE: $(( CURRENT - 1 )) is the last plugin argument completely passed, i.e. that which
+        # has a space after them. This is to avoid removing plugins partially passed, which makes
+        # the completion not add a space after the completed plugin.
+        local -a args
+        args=(${words[4,$(( CURRENT - 1))]})
+        valid_plugins=(${valid_plugins:|args})
+
+        _describe 'plugin' valid_plugins ;;
+    esac
+  fi
+
+  return 0
+}
+
+# If run from a script, do not set the completion function
+if (( ${+functions[compdef]} )); then
+  compdef _omz omz
+fi
+
 ## Utility functions
 
 function _omz::confirm {
@@ -844,5 +942,3 @@ function _omz::version {
     printf "%s (%s)\n" "$version" "$commit"
   )
 }
-
-omz "$@"

lib/diagnostics.zsh

@@ -1,4 +1,4 @@
-# omz_diagnostic_dump
+# diagnostics.zsh
 #
 # Diagnostic and debugging support for oh-my-zsh
 
@@ -351,5 +351,3 @@ function _omz_diag_dump_os_specific_version() {
   done
 }
 
-
-omz_diagnostic_dump "$@"

lib/vcs_info.zsh

@@ -38,16 +38,47 @@
 # due to malicious input as a consequence of CVE-2021-45444, which affects
 # zsh versions from 5.0.3 to 5.8.
 #
-autoload -Uz +X regexp-replace VCS_INFO_formats 2>/dev/null || return 0
+# The patch is applied when VCS_INFO_formats is first called, since loading
+# and patching it on every startup costs time in shells that never use
+# vcs_info. `autoload` doesn't replace an already defined function, so this
+# wrapper survives a later `autoload -Uz vcs_info` in a theme or .zshrc, and
+# vcs_info's own `autoload -Uz VCS_INFO_formats`.
+if (( $+functions[VCS_INFO_formats] )); then
+  () {
+    autoload -Uz +X regexp-replace 2>/dev/null || return 1
 
-# We use $tmp here because it's already a local variable in VCS_INFO_formats
-typeset PATCH='for tmp (base base-name branch misc revision subdir) hook_com[$tmp]="${hook_com[$tmp]//\%/%%}"'
-# Unique string to avoid reapplying the patch if this code gets called twice
-typeset PATCH_ID=vcs_info-patch-9b9840f2-91e5-4471-af84-9e9a0dc68c1b
-# Only patch the VCS_INFO_formats function if not already patched
-if [[ "$functions[VCS_INFO_formats]" != *$PATCH_ID* ]]; then
-  regexp-replace 'functions[VCS_INFO_formats]' \
-    "VCS_INFO_hook 'post-backend'" \
-    ': ${PATCH_ID}; ${PATCH}; ${MATCH}'
+    # We use $tmp here because it's already a local variable in VCS_INFO_formats
+    local PATCH='for tmp (base base-name branch misc revision subdir) hook_com[$tmp]="${hook_com[$tmp]//\%/%%}"'
+    # Unique string to avoid reapplying the patch if this code gets called twice
+    local PATCH_ID=vcs_info-patch-9b9840f2-91e5-4471-af84-9e9a0dc68c1b
+    # Only patch the VCS_INFO_formats function if not already patched
+    if [[ "$functions[VCS_INFO_formats]" != *$PATCH_ID* ]]; then
+      regexp-replace 'functions[VCS_INFO_formats]' \
+        "VCS_INFO_hook 'post-backend'" \
+        ': ${PATCH_ID}; ${PATCH}; ${MATCH}'
+    fi
+  }
+else
+  function VCS_INFO_formats {
+    local loaded_function="$(
+      unfunction VCS_INFO_formats 2>/dev/null
+      autoload -Uz +X VCS_INFO_formats 2>/dev/null || return 1
+      print -r -- "$functions[VCS_INFO_formats]"
+    )" || return 1
+    functions[VCS_INFO_formats]="$loaded_function"
+    autoload -Uz +X regexp-replace 2>/dev/null || return 1
+
+    # We use $tmp here because it's already a local variable in VCS_INFO_formats
+    local PATCH='for tmp (base base-name branch misc revision subdir) hook_com[$tmp]="${hook_com[$tmp]//\%/%%}"'
+    # Unique string to avoid reapplying the patch if this code gets called twice
+    local PATCH_ID=vcs_info-patch-9b9840f2-91e5-4471-af84-9e9a0dc68c1b
+    # Only patch the VCS_INFO_formats function if not already patched
+    if [[ "$functions[VCS_INFO_formats]" != *$PATCH_ID* ]]; then
+      regexp-replace 'functions[VCS_INFO_formats]' \
+        "VCS_INFO_hook 'post-backend'" \
+        ': ${PATCH_ID}; ${PATCH}; ${MATCH}'
+    fi
+
+    VCS_INFO_formats "$@"
+  }
 fi
-unset PATCH PATCH_ID

oh-my-zsh.sh

@@ -78,9 +78,6 @@ fpath=($ZSH/{functions,completions} $ZSH_CUSTOM/{functions,completions} $fpath)
 # Load all stock functions (from $fpath files) called below.
 autoload -U compaudit compinit zrecompile
 
-# The omz CLI and omz_diagnostic_dump are loaded on first use
-autoload -Uz omz omz_diagnostic_dump
-
 is_plugin() {
   local base_dir=$1
   local name=$2