From f0942e6dda6c4bc0711a58efce71d4b11eea7b5b Mon Sep 17 00:00:00 2001
From: Indrajit Raychaudhuri <irc@indrajit.com>
Date: Mon, 30 Nov 2020 23:55:27 -0600
Subject: [PATCH] python: Optimize completion for loading lazily on demand

Move `pip` to separate completion definition delegating to pip to do
all the completion work.

Since the completion is loaded lazily on-demand, it avoids the
performance overhead during Zsh initialization.

Implementation note:
The helper function `_pip_completion` implementation is based on the
official pip completion function (which can be generated with
`pip completion --zsh`) adhering to the newer compsys style.

See: 'man zshcompsys' for more details.
---
 modules/python/README.md      |  5 ++++-
 modules/python/functions/_pip | 19 +++++++++++++++++++
 modules/python/init.zsh       | 32 +-------------------------------
 3 files changed, 24 insertions(+), 32 deletions(-)
 create mode 100644 modules/python/functions/_pip

diff --git a/modules/python/README.md b/modules/python/README.md
index fa5140a..3c56612 100644
--- a/modules/python/README.md
+++ b/modules/python/README.md
@@ -2,6 +2,9 @@
 
 Enables local Python and local Python package installation.
 
+This module must be loaded _before_ the _`completion`_ module so that the
+provided completion definitions are loaded.
+
 ## Settings
 
 This module supports virtual environments from conda and
@@ -143,7 +146,7 @@ Then add `$python_info[virtualenv]` to `$PROMPT` or `$RPROMPT` and call
 
 Similarly, you can use `:prezto:module:python:info:version:format` with `%v` for
 the version and add `$python_info[version]` to your prompt for the current
-python version/
+python version.
 
 ## Authors
 
diff --git a/modules/python/functions/_pip b/modules/python/functions/_pip
new file mode 100644
index 0000000..4e0ffb8
--- /dev/null
+++ b/modules/python/functions/_pip
@@ -0,0 +1,19 @@
+#compdef -P pip[0-9.]#
+#autoload
+
+#
+# Pip completion, delegating to pip to do all the completion work.
+#
+# Authors:
+#   Indrajit Raychaudhuri <irc@indrajit.com>
+#
+
+if (( $+commands[$words[1]] )); then
+
+  function _pip_completion {
+    compadd -- $( COMP_WORDS="$words[*]" COMP_CWORD=$(( CURRENT - 1 )) \
+                  PIP_AUTO_COMPLETE=1 $words[1] 2>/dev/null )
+  }
+  _pip_completion "$@"
+
+fi
diff --git a/modules/python/init.zsh b/modules/python/init.zsh
index 29fc8dc..0682b43 100644
--- a/modules/python/init.zsh
+++ b/modules/python/init.zsh
@@ -5,6 +5,7 @@
 #   Sorin Ionescu <sorin.ionescu@gmail.com>
 #   Sebastian Wiesner <lunaryorn@googlemail.com>
 #   Patrick Bos <egpbos@gmail.com>
+#   Indrajit Raychaudhuri <irc@indrajit.com>
 #
 
 # Load dependencies.
@@ -133,37 +134,6 @@ if (( $+VIRTUALENVWRAPPER_VIRTUALENV || $+commands[virtualenv] )) \
   unset pyenv_plugins
 fi
 
-# Load PIP completion.
-# Detect and use one available from among 'pip', 'pip2', 'pip3' variants
-if [[ -n "$PYENV_ROOT" ]]; then
-  for pip in pip{,2,3}; do
-    pip_command="$(pyenv which "$pip" 2> /dev/null)"
-    [[ -n "$pip_command" ]] && break
-  done
-  unset pip
-else
-  pip_command="$commands[(i)pip(|[23])]"
-fi
-if [[ -n "$pip_command" ]]; then
-  cache_file="${XDG_CACHE_HOME:-$HOME/.cache}/prezto/pip-cache.zsh"
-
-  if [[ "$pip_command" -nt "$cache_file" \
-        || "${ZDOTDIR:-$HOME}/.zpreztorc" -nt "$cache_file" \
-        || ! -s "$cache_file" ]]; then
-    mkdir -p "$cache_file:h"
-    # pip is slow; cache its output. And also support 'pip2', 'pip3' variants
-    "$pip_command" completion --zsh \
-      | sed -e "s/\(compctl -K [-_[:alnum:]]* pip\).*/\1{,2,3}{,.{0..9}}/" \
-      >! "$cache_file" \
-      2> /dev/null
-  fi
-
-  source "$cache_file"
-
-  unset cache_file
-fi
-unset pip_command
-
 # Load conda into the shell session, if requested.
 zstyle -T ':prezto:module:python' conda-init
 if (( $? && $+commands[conda] )); then