Add Poetry support to virtualenv

config/p10k-classic.zsh

@@ -944,6 +944,8 @@
   # If set to "false", won't show virtualenv if pyenv is already shown.
   # If set to "if-different", won't show virtualenv if it's the same as pyenv.
   typeset -g POWERLEVEL9K_VIRTUALENV_SHOW_WITH_PYENV=false
+  # Check for virtual environments managed by poetry
+  typeset -g POWERLEVEL9K_VIRTUALENV_ENABLE_POETRY=true
   # Separate environment name from Python version only with a space.
   typeset -g POWERLEVEL9K_VIRTUALENV_{LEFT,RIGHT}_DELIMITER=
   # Custom icon.

config/p10k-lean-8colors.zsh

@@ -922,6 +922,8 @@
   # If set to "false", won't show virtualenv if pyenv is already shown.
   # If set to "if-different", won't show virtualenv if it's the same as pyenv.
   typeset -g POWERLEVEL9K_VIRTUALENV_SHOW_WITH_PYENV=false
+  # Check for virtual environments managed by poetry
+  typeset -g POWERLEVEL9K_VIRTUALENV_ENABLE_POETRY=true
   # Separate environment name from Python version only with a space.
   typeset -g POWERLEVEL9K_VIRTUALENV_{LEFT,RIGHT}_DELIMITER=
   # Custom icon.

config/p10k-lean.zsh

@@ -918,6 +918,8 @@
   # If set to "false", won't show virtualenv if pyenv is already shown.
   # If set to "if-different", won't show virtualenv if it's the same as pyenv.
   typeset -g POWERLEVEL9K_VIRTUALENV_SHOW_WITH_PYENV=false
+  # Check for virtual environments managed by poetry
+  typeset -g POWERLEVEL9K_VIRTUALENV_ENABLE_POETRY=true
   # Separate environment name from Python version only with a space.
   typeset -g POWERLEVEL9K_VIRTUALENV_{LEFT,RIGHT}_DELIMITER=
   # Custom icon.

config/p10k-pure.zsh

@@ -93,6 +93,8 @@
   # Don't show Python version.
   typeset -g POWERLEVEL9K_VIRTUALENV_SHOW_PYTHON_VERSION=false
   typeset -g POWERLEVEL9K_VIRTUALENV_{LEFT,RIGHT}_DELIMITER=
+  # Check for virtual environments managed by poetry
+  typeset -g POWERLEVEL9K_VIRTUALENV_ENABLE_POETRY=true
 
   # Blue current directory.
   typeset -g POWERLEVEL9K_DIR_FOREGROUND=$blue

config/p10k-rainbow.zsh

@@ -988,6 +988,8 @@
   # If set to "false", won't show virtualenv if pyenv is already shown.
   # If set to "if-different", won't show virtualenv if it's the same as pyenv.
   typeset -g POWERLEVEL9K_VIRTUALENV_SHOW_WITH_PYENV=false
+  # Check for virtual environments managed by poetry
+  typeset -g POWERLEVEL9K_VIRTUALENV_ENABLE_POETRY=true
   # Separate environment name from Python version only with a space.
   typeset -g POWERLEVEL9K_VIRTUALENV_{LEFT,RIGHT}_DELIMITER=
   # Custom icon.

internal/p10k.zsh

@@ -2403,6 +2403,7 @@ _p9k_prompt_load_sync() {
 # The first argument says whether to capture stderr (1) or ignore it (0).
 # The second argument can be empty or a file. If it's a file, the
 # output of the command is presumed to potentially depend on it.
+# Returns output in $_p9k__ret.
 function _p9k_cached_cmd() {
   local cmd=$commands[$3]
   [[ -n $cmd ]] || return
@@ -4280,24 +4281,56 @@ function instant_prompt_chezmoi_shell() {
 # Virtualenv: current working virtualenv
 # More information on virtualenv (Python):
 # https://virtualenv.pypa.io/en/latest/
+
+_virtualenv_VIRTUAL_ENV() {
+  local n
+  # VIRTUAL_ENV is the on-disk path, VIRTUAL_ENV_PROMPT is `(name_of_virtualenv) `.
+  # Use custom name if it was set (python -m venv -p "foo" .venv)
+  if [[ $VIRTUAL_ENV_PROMPT == '('?*') ' && $VIRTUAL_ENV_PROMPT != "($n) " ]]; then
+    n=$VIRTUAL_ENV_PROMPT[2,-3]
+  # Use parent directory name if virtualenv name is generic (eg .venv)
+  elif [[ $v == $~_POWERLEVEL9K_VIRTUALENV_GENERIC_NAMES ]]; then
+    n=${VIRTUAL_ENV:h:t}
+  # Otherwise use the virtualenv name as-is
+  else
+    n=${VIRTUAL_ENV:t}
+  fi
+  echo "$n"
+}
+_virtualenv_poetry() {
+  local idx=$1
+  local dir=${_p9k__parent_dirs[idx]}
+  local pyproject="$dir/pyproject.toml"
+  _p9k_cached_cmd 0 '' poetry -C "$dir" version
+  # Return the first word only, eg the value of pyproject.toml's `poetry.name`
+  echo "${_p9k__ret%% *}"
+}
 prompt_virtualenv() {
   local msg=''
   if (( _POWERLEVEL9K_VIRTUALENV_SHOW_PYTHON_VERSION )) && _p9k_python_version; then
     msg="${_p9k__ret//\%/%%} "
   fi
-  local v=${VIRTUAL_ENV:t}
-  if [[ $VIRTUAL_ENV_PROMPT == '('?*') ' && $VIRTUAL_ENV_PROMPT != "($v) " ]]; then
-    v=$VIRTUAL_ENV_PROMPT[2,-3]
-  elif [[ $v == $~_POWERLEVEL9K_VIRTUALENV_GENERIC_NAMES ]]; then
-    v=${VIRTUAL_ENV:h:t}
+  # Determine virtualenv name with a few strategies
+  local n=''
+  if [[ -n $VIRTUAL_ENV ]]; then
+    n=$(_virtualenv_VIRTUAL_ENV)
+  else
+    if [[ $POWERLEVEL9K_VIRTUALENV_ENABLE_POETRY == true ]]; then
+      local start end
+      _p9k_upglob pyproject.toml
+      local idx=$?
+      if (( idx > 0 )); then
+        n=$(_virtualenv_poetry $idx)
+      fi
+    fi
   fi
-  msg+="$_POWERLEVEL9K_VIRTUALENV_LEFT_DELIMITER${v//\%/%%}$_POWERLEVEL9K_VIRTUALENV_RIGHT_DELIMITER"
+  msg+="$_POWERLEVEL9K_VIRTUALENV_LEFT_DELIMITER${n//\%/%%}$_POWERLEVEL9K_VIRTUALENV_RIGHT_DELIMITER"
   case $_POWERLEVEL9K_VIRTUALENV_SHOW_WITH_PYENV in
     false)
       _p9k_prompt_segment "$0" "blue" "$_p9k_color1" 'PYTHON_ICON' 0 '${(M)${#P9K_PYENV_PYTHON_VERSION}:#0}' "$msg"
     ;;
     if-different)
-      _p9k_escape $v
+      _p9k_escape $n
       _p9k_prompt_segment "$0" "blue" "$_p9k_color1" 'PYTHON_ICON' 0 '${${:-'$_p9k__ret'}:#$_p9k__pyenv_version}' "$msg"
     ;;
     *)
@@ -4307,7 +4340,8 @@ prompt_virtualenv() {
 }
 
 _p9k_prompt_virtualenv_init() {
-  typeset -g "_p9k__segment_cond_${_p9k__prompt_side}[_p9k__segment_index]"='$VIRTUAL_ENV'
+  # Init if VIRTUAL_ENV is non-null or poetry is installed
+  typeset -g "_p9k__segment_cond_${_p9k__prompt_side}[_p9k__segment_index]"='${VIRTUAL_ENV:-${commands[poetry]:-${${+functions[poetry]}:#0}}}'
 }
 
 # _p9k_read_pyenv_like_version_file <filepath> [prefix]