2 files changed, 50 insertions, 17 deletions
diff --git a/plugins/pyenv/README.md b/plugins/pyenv/README.md
index d063b55b9..b9ee937b7 100644
--- a/plugins/pyenv/README.md
+++ b/plugins/pyenv/README.md
@@ -1,4 +1,4 @@
-# pyenv 
+# pyenv
 
 This plugin looks for [pyenv](https://github.com/pyenv/pyenv), a Simple Python version
 management system, and loads it if it's found. It also loads pyenv-virtualenv, a pyenv
@@ -10,6 +10,14 @@ To use it, add `pyenv` to the plugins array in your zshrc file:
 plugins=(... pyenv)
 ```
 
+## Settings
+
+- `ZSH_PYENV_QUIET`: if set to `true`, the plugin will not print any messages if it
+  finds that `pyenv` is not properly configured.
+
+- `ZSH_PYENV_VIRTUALENV`: if set to `false`, the plugin will not load pyenv-virtualenv
+  when it finds it.
+
 ## Functions
 
 - `pyenv_prompt_info`: displays the Python version in use by pyenv; or the global Python
diff --git a/plugins/pyenv/pyenv.plugin.zsh b/plugins/pyenv/pyenv.plugin.zsh
index 813f64b42..d91b5daa7 100644
--- a/plugins/pyenv/pyenv.plugin.zsh
+++ b/plugins/pyenv/pyenv.plugin.zsh
@@ -1,3 +1,24 @@
+pyenv_config_warning() {
+  [[ "$ZSH_PYENV_QUIET" != true ]] || return 0
+
+  local reason="$1"
+  local pyenv_root="${PYENV_ROOT/#$HOME/\$HOME}"
+  cat >&2 <<EOF
+Found pyenv, but it is badly configured ($reason). pyenv might not
+work correctly for non-interactive shells (for example, when run from a script).
+${(%):-"%B%F{yellow}"}
+To fix this message, add these lines to the '.profile' and '.zprofile' files
+in your home directory:
+${(%):-"%f"}
+export PYENV_ROOT="$pyenv_root"
+export PATH="\$PYENV_ROOT/bin:\$PATH"
+eval "\$(pyenv init --path)"
+${(%):-"%F{yellow}"}
+You'll need to restart your user session for the changes to take effect.${(%):-%b%f}
+For more information go to https://github.com/pyenv/pyenv/#installation.
+EOF
+}
+
 # This plugin loads pyenv into the current shell and provides prompt info via
 # the 'pyenv_prompt_info' function. Also loads pyenv-virtualenv if available.
 
@@ -30,31 +51,34 @@ if [[ $FOUND_PYENV -ne 1 ]]; then
 
   # If we found pyenv, load it but show a caveat about non-interactive shells
   if [[ $FOUND_PYENV -eq 1 ]]; then
-    cat <<EOF
-Found pyenv, but it is badly configured. pyenv might not work for
-non-interactive shells (for example, when run from a script).
-${bold_color}
-To fix this message, add these lines to the '.profile' and '.zprofile' files
-in your home directory:
-
-export PYENV_ROOT="${dir/#$HOME/\$HOME}"
-export PATH="\$PYENV_ROOT/bin:\$PATH"
-eval "\$(pyenv init --path)"
-${reset_color}
-For more info go to https://github.com/pyenv/pyenv/#installation.
-EOF
-
     # Configuring in .zshrc only makes pyenv available for interactive shells
-    export PYENV_ROOT=$dir
+    export PYENV_ROOT="$dir"
     export PATH="$PYENV_ROOT/bin:$PATH"
     eval "$(pyenv init --path)"
+
+    # Show warning due to bad pyenv configuration
+    pyenv_config_warning 'pyenv command not found in $PATH'
   fi
 fi
 
 if [[ $FOUND_PYENV -eq 1 ]]; then
+  if [[ -z "$PYENV_ROOT" ]]; then
+    # This is only for backwards compatibility with users that previously relied
+    # on this plugin exporting it. pyenv itself does not require it to be exported
+    export PYENV_ROOT="$(pyenv root)"
+  fi
+
+  # Add pyenv shims to $PATH if not already added
+  if [[ -z "${path[(Re)$(pyenv root)/shims]}" ]]; then
+    eval "$(pyenv init --path)"
+    pyenv_config_warning 'missing pyenv shims in $PATH'
+  fi
+
+  # Load pyenv
   eval "$(pyenv init - --no-rehash zsh)"
 
-  if (( ${+commands[pyenv-virtualenv-init]} )); then
+  # If pyenv-virtualenv exists, load it
+  if [[ -d "$(pyenv root)/plugins/pyenv-virtualenv" && "$ZSH_PYENV_VIRTUALENV" != false ]]; then
     eval "$(pyenv virtualenv-init - zsh)"
   fi
 
@@ -69,3 +93,4 @@ else
 fi
 
 unset FOUND_PYENV pyenvdirs dir
+unfunction pyenv_config_warning