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

2 files changed, 59 insertions, 12 deletions

diff --git a/plugins/python/README.md b/plugins/python/README.md
index 8014992de..490cf88f1 100644
--- a/plugins/python/README.md
+++ b/plugins/python/README.md
@@ -24,16 +24,45 @@ plugins=(... python)
 The plugin provides three utilities to manage Python 3.3+ [venv](https://docs.python.org/3/library/venv.html)
 virtual environments:
 
-- `mkv [name]`: make a new virtual environment called `name` (default: if set `$PYTHON_VENV_NAME`, else
-  `venv`) in the current directory.
+- `mkv [name]`: make a new virtual environment called `name` in the current directory.
+  **Default**: `$PYTHON_VENV_NAME` if set, otherwise `venv`.
 
-- `vrun [name]`: Activate the virtual environment called `name` (default: if set `$PYTHON_VENV_NAME`, else
-  `venv`) in the current directory.
+- `vrun [name]`: activate the virtual environment called `name` in the current directory.
+  **Default**: the first existing in `$PYTHON_VENV_NAMES`.
 
-- `auto_vrun`: Automatically activate the venv virtual environment when entering a directory containing
+- `auto_vrun`: automatically activate the venv virtual environment when entering a directory containing
   `<venv-name>/bin/activate`, and automatically deactivate it when navigating out of it (keeps venv activated
   in subdirectories).
-  - To enable the feature, set `export PYTHON_AUTO_VRUN=true` before sourcing oh-my-zsh.
-  - Plugin activates first virtual environment in lexicographic order whose name begins with `<venv-name>`.
+  - To enable the feature, set `PYTHON_AUTO_VRUN=true` before sourcing oh-my-zsh.
+  - The plugin activates the first existing virtual environment, in order, appearing in `$PYTON_VENV_NAMES`.
     The default virtual environment name is `venv`. To use a different name, set
-    `export PYTHON_VENV_NAME=<venv-name>`. For example: `export PYTHON_VENV_NAME=".venv"`
+    `PYTHON_VENV_NAME=<venv-name>`. For example: `PYTHON_VENV_NAME=".venv"`
+
+### Settings
+
+You can set these variables in your `.zshrc` file, before Oh My Zsh is sourced.
+For example:
+
+```sh
+PYTHON_VENV_NAME=".venv"
+PYTHON_VENV_NAMES=($PYTHON_VENV_NAME venv)
+...
+plugins=(... python)
+source "$ZSH/oh-my-zsh.sh"
+```
+
+
+## `$PYTHON_VENV_NAME`
+
+**Default**: `venv`.
+
+Preferred name for virtual environments, for example when creating via `mkv`.
+
+## `$PYTHON_VENV_NAMES`
+
+**Default**: `$PYTHON_VENV_NAME venv .venv`.
+
+Array of virtual environment names to be checked, in order, by `vrun` and `auto_vrun`.
+This means these functions will load the first existing virtual environment in this list.
+Duplicate names are ignored.
+
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