From 3ab569ff786c324a0c597c1d1cff44491565d0f8 Mon Sep 17 00:00:00 2001
From: Kaleb Elwert <belak@coded.io>
Date: Mon, 4 Dec 2017 15:08:58 -0800
Subject: [PATCH] python-info: add support for pyenv and only do extra work if
 the prompt needs it

---
 modules/python/functions/python-info | 38 ++++++++++++++++++----------
 1 file changed, 24 insertions(+), 14 deletions(-)

diff --git a/modules/python/functions/python-info b/modules/python/functions/python-info
index e5a8339..144a5b6 100644
--- a/modules/python/functions/python-info
+++ b/modules/python/functions/python-info
@@ -13,29 +13,39 @@ local virtualenv_format
 local virtualenv_formatted
 local version_format
 local version_formatted
+local version
 
 # Clean up previous $python_info.
 unset python_info
 typeset -gA python_info
 
+# Grab the styling we might have to do
+zstyle -s ':prezto:module:python:info:virtualenv' format 'virtualenv_format'
+zstyle -s ':prezto:module:python:info:version' format 'version_format'
+
 # Format virtualenv.
-if [[ -n "$VIRTUAL_ENV" ]]; then
-  zstyle -s ':prezto:module:python:info:virtualenv' format 'virtualenv_format'
-  zformat -f virtualenv_formatted "$virtualenv_format" "v:${VIRTUAL_ENV:t}"
-  python_info[virtualenv]="$virtualenv_formatted"
-fi
+if [[ -n "$virtualenv_format" ]]; then
+  if [[ -n "$VIRTUAL_ENV" ]]; then
+    zformat -f virtualenv_formatted "$virtualenv_format" "v:${VIRTUAL_ENV:t}"
+    python_info[virtualenv]="$virtualenv_formatted"
+  fi
 
-# Do the same for Conda virtual environments
-if [[ -n "$CONDA_DEFAULT_ENV" ]]; then
-  zstyle -s ':prezto:module:python:info:virtualenv' format 'virtualenv_format'
-  zformat -f virtualenv_formatted "$virtualenv_format" "v:${CONDA_DEFAULT_ENV:t}"
-  python_info[virtualenv]="$virtualenv_formatted"
+  # Do the same for Conda virtual environments
+  if [[ -n "$CONDA_DEFAULT_ENV" ]]; then
+    zformat -f virtualenv_formatted "$virtualenv_format" "v:${CONDA_DEFAULT_ENV:t}"
+    python_info[virtualenv]="$virtualenv_formatted"
+  fi
 fi
 
-if (( $+commands[python] )); then
-  zstyle -s ':prezto:module:python:info:version' format 'version_format'
-  if [[ -n "$version_format" ]]; then
-    zformat -f version_formatted "$version_format" "v:${$(python3 --version)#Python }"
+if [[ -n "$version_format" ]]; then
+  if (( $+commands[pyenv] )); then
+    version="${"$(pyenv version)"%% *}"
+  elif (( $+commands[python] )); then
+    version="${$(python3 --version)#Python }"
+  fi
+
+  if [[ -n "$version" && "$version" != "system" ]]; then
+    zformat -f version_formatted "$version_format" "v:$version"
     python_info[version]="$version_formatted"
   fi
 fi