From ea07fe6abc9596a9b19314a0d0cf37693177a27d Mon Sep 17 00:00:00 2001
From: Indrajit Raychaudhuri <irc@indrajit.com>
Date: Mon, 24 May 2021 17:31:17 -0500
Subject: [PATCH] python: Update pyenv initialization flow per pyenv 2.0

In pyenv 2.0 onwards, it is not enough anymore to initialize pyenv in
shell by just calling `pyenv init -`. We also need to update `path` to
include pyenv shims by calling `pyenv init --path`.

Also, honor `$PYENV_ROOT` if set but, no need to set it explicitly if
not set. Instead, let the initialization script take care of that.
---
 modules/python/README.md |  4 +++-
 modules/python/init.zsh  | 23 ++++++++++++++++-------
 2 files changed, 19 insertions(+), 8 deletions(-)

diff --git a/modules/python/README.md b/modules/python/README.md
index 3c56612..1695fa0 100644
--- a/modules/python/README.md
+++ b/modules/python/README.md
@@ -38,7 +38,9 @@ execution of `pyenv`.
 Install Python versions with `pyenv install` into _`~/.pyenv/versions`_.
 
 This will be loaded automatically if pyenv is installed to `$PYENV_ROOT`,
-_`~/.pyenv`_, or if the `pyenv` command is on the path.
+_`~/.pyenv`_, or if the `pyenv` command is on the path. By default, `PYENV_ROOT`
+is set to _`$HOME/.pyenv`_. You can set it to an alternate location and export
+it in [_`${ZDOTDIR:-$HOME}/.zshenv`_][6].
 
 ## Local Package Installation
 
diff --git a/modules/python/init.zsh b/modules/python/init.zsh
index 0682b43..ed68990 100644
--- a/modules/python/init.zsh
+++ b/modules/python/init.zsh
@@ -11,14 +11,21 @@
 # Load dependencies.
 pmodload 'helper'
 
-# Load manually installed pyenv into the path
-if [[ -s "${PYENV_ROOT:=$HOME/.pyenv}/bin/pyenv" ]]; then
-  path=("$PYENV_ROOT/bin" $path)
-  eval "$(pyenv init - --no-rehash zsh)"
+# Load manually installed or package manager installed pyenv into the shell
+# session.
+if [[ -s "${local_pyenv::=${PYENV_ROOT:-$HOME/.pyenv}/bin/pyenv}" ]] \
+      || (( $+commands[pyenv] )); then
+
+  # Ensure manually installed pyenv is added to path when present.
+  [[ -s $local_pyenv ]] && path=($local_pyenv:h $path)
+
+  # pyenv 2+ requires shims to be added to path before being initialized.
+  autoload -Uz is-at-least
+  if is-at-least 2 ${"$(pyenv --version 2>&1)"[(w)2]}; then
+    eval "$(pyenv init --path zsh)"
+  fi
 
-# Load pyenv into the current python session
-elif (( $+commands[pyenv] )); then
-  eval "$(pyenv init - --no-rehash zsh)"
+  eval "$(pyenv init - zsh)"
 
 # Prepend PEP 370 per user site packages directory, which defaults to
 # ~/Library/Python on macOS and ~/.local elsewhere, to PATH. The
@@ -34,6 +41,8 @@ else
   fi
 fi
 
+unset local_pyenv
+
 # Return if requirements are not found.
 if (( ! $#commands[(i)python[23]#] && ! $+functions[pyenv] )); then
   return 1