From 12cd3b3e399d39b2b458fdd8f1f6286250253476 Mon Sep 17 00:00:00 2001
From: Robert Crews <rcrews@mac.com>
Date: Sun, 3 Mar 2024 04:04:19 -0800
Subject: [PATCH] feat(python): add auto venv activation (#12248)

Co-authored-by: Carlo Sala <carlosalag@protonmail.com>
---
 plugins/python/README.md         | 16 +++++++++++++---
 plugins/python/python.plugin.zsh | 23 +++++++++++++++++++----
 2 files changed, 32 insertions(+), 7 deletions(-)

diff --git a/plugins/python/README.md b/plugins/python/README.md
index 7bf1b34ac..c99697b22 100644
--- a/plugins/python/README.md
+++ b/plugins/python/README.md
@@ -22,8 +22,18 @@ plugins=(... python)
 
 ## Virtual environments
 
-The plugin provides two utilities to manage Python venvs:
+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: `venv`) in current directory.
+- `mkv [name]`: make a new virtual environment called `name` (default: if set `$PYTHON_VENV_NAME`, else
+  `venv`) in the current directory.
 
-- `vrun [name]`: activate virtual environment called `name` (default: `venv`) in current directory.
+- `vrun [name]`: Activate the virtual environment called `name` (default: if set `$PYTHON_VENV_NAME`, else
+  `venv`) in the current directory.
+
+- `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 (including
+  subdirectories!).
+  - To enable the feature, set `export PYTHON_AUTO_VRUN=true` before sourcing oh-my-zsh.
+  - 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"`
diff --git a/plugins/python/python.plugin.zsh b/plugins/python/python.plugin.zsh
index 77d4bf425..f6ea85027 100644
--- a/plugins/python/python.plugin.zsh
+++ b/plugins/python/python.plugin.zsh
@@ -51,11 +51,12 @@ alias pyserver="python3 -m http.server"
 
 
 ## venv utilities
+: ${PYTHON_VENV_NAME:=venv}
 
 # Activate a the python virtual environment specified.
-# If none specified, use 'venv'.
+# If none specified, use $PYTHON_VENV_NAME, else 'venv'.
 function vrun() {
-  local name="${1:-venv}"
+  local name="${1:-$PYTHON_VENV_NAME}"
   local venvpath="${name:P}"
 
   if [[ ! -d "$venvpath" ]]; then
@@ -72,12 +73,26 @@ function vrun() {
   echo "Activated virtual environment ${name}"
 }
 
-# Create a new virtual environment, with default name 'venv'.
+# Create a new virtual environment using the specified name.
+# If none specfied, use $PYTHON_VENV_NAME
 function mkv() {
-  local name="${1:-venv}"
+  local name="${1:-$PYTHON_VENV_NAME}"
   local venvpath="${name:P}"
 
   python3 -m venv "${name}" || return
   echo >&2 "Created venv in '${venvpath}'"
   vrun "${name}"
 }
+
+if [[ "$PYTHON_AUTO_VRUN" == "true" ]]; then
+  # Automatically activate venv when changing dir
+  auto_vrun() {
+    if [[ -f "${PYTHON_VENV_NAME}/bin/activate" ]]; then
+      source "${PYTHON_VENV_NAME}/bin/activate" > /dev/null 2>&1
+    else
+      (( $+functions[deactivate] )) && deactivate > /dev/null 2>&1
+    fi
+  }
+  add-zsh-hook chpwd auto_vrun
+  auto_vrun
+fi