feat(python): support multiple venvs via `$PYTHON_VENV_NAMES` (#12932)

1 files changed, 22 insertions, 4 deletions

diff --git a/plugins/python/python.plugin.zsh b/plugins/python/python.plugin.zsh
index 63733e1de..2b139ddf0 100644
--- a/plugins/python/python.plugin.zsh
+++ b/plugins/python/python.plugin.zsh
@@ -47,12 +47,29 @@ alias pygrep='grep -nr --include="*.py"'
 alias pyserver="python3 -m http.server"
 
 
-## venv utilities
+## venv settings
 : ${PYTHON_VENV_NAME:=venv}
 
+# Array of possible virtual environment names to look for, in order
+# -U for removing duplicates
+typeset -gaU PYTHON_VENV_NAMES
+[[ -n "$PYTHON_VENV_NAMES" ]] || PYTHON_VENV_NAMES=($PYTHON_VENV_NAME venv .venv)
+
 # Activate a the python virtual environment specified.
-# If none specified, use $PYTHON_VENV_NAME, else 'venv'.
+# If none specified, use the first existing in $PYTHON_VENV_NAMES.
 function vrun() {
+  if [[ -z "$1" ]]; then
+    local name
+    for name in $PYTHON_VENV_NAMES; do
+      local venvpath="${name:P}"
+      if [[ -d "$venvpath" ]]; then
+        vrun "$name"
+        return $?
+      fi
+    done
+    echo >&2 "Error: no virtual environment found in current directory"
+  fi
+
   local name="${1:-$PYTHON_VENV_NAME}"
   local venvpath="${name:P}"
 
@@ -91,10 +108,11 @@ if [[ "$PYTHON_AUTO_VRUN" == "true" ]]; then
     fi
 
     if [[ $PWD != ${VIRTUAL_ENV:h} ]]; then
-      for _file in "${PYTHON_VENV_NAME}"*/bin/activate(N.); do
+      local file
+      for file in "${^PYTHON_VENV_NAMES[@]}"/bin/activate(N.); do
         # make sure we're not in a venv already
         (( $+functions[deactivate] )) && deactivate > /dev/null 2>&1
-        source $_file > /dev/null 2>&1
+        source $file > /dev/null 2>&1
         break
       done
     fi