From 31f2025e0fa963788655fe197e0179c47588b175 Mon Sep 17 00:00:00 2001
From: Kate Sullivan <75387802+katesullivan@users.noreply.github.com>
Date: Fri, 12 Apr 2024 02:24:54 -0500
Subject: [PATCH] feat(poetry-env): support changing between two venv dirs
 (#12346)

Co-authored-by: Carlo Sala <carlosalag@protonmail.com>
---
 plugins/poetry-env/poetry-env.plugin.zsh | 38 ++++++++++++------------
 1 file changed, 19 insertions(+), 19 deletions(-)

diff --git a/plugins/poetry-env/poetry-env.plugin.zsh b/plugins/poetry-env/poetry-env.plugin.zsh
index 97ca52492..f98c177b4 100644
--- a/plugins/poetry-env/poetry-env.plugin.zsh
+++ b/plugins/poetry-env/poetry-env.plugin.zsh
@@ -1,27 +1,27 @@
-# Automatic poetry environment activation/deactivation
 _togglePoetryShell() {
-  # deactivate environment if pyproject.toml doesn't exist and not in a subdir
-  if [[ ! -f "$PWD/pyproject.toml" ]] ; then
-    if [[ "$poetry_active" == 1 ]]; then
-      if [[ "$PWD" != "$poetry_dir"* ]]; then
-        export poetry_active=0
-        deactivate
-        return
-      fi
-    fi
+  # Determine if currently in a Poetry-managed directory
+  local in_poetry_dir=0
+  if [[ -f "$PWD/pyproject.toml" ]] && grep -q 'tool.poetry' "$PWD/pyproject.toml"; then
+    in_poetry_dir=1
   fi
 
-  # activate the environment if pyproject.toml exists
-  if [[ "$poetry_active" != 1 ]]; then
-    if [[ -f "$PWD/pyproject.toml" ]]; then
-      if grep -q 'tool.poetry' "$PWD/pyproject.toml" && venv_dir=$(poetry env info --path); then
-        export poetry_active=1
-        export poetry_dir="$PWD"
-        source "${venv_dir}/bin/activate"
-      fi
+  # Deactivate the current environment if moving out of a Poetry directory or into a different Poetry directory
+  if [[ $poetry_active -eq 1 ]] && { [[ $in_poetry_dir -eq 0 ]] || [[ "$PWD" != "$poetry_dir"* ]]; }; then
+    export poetry_active=0
+    unset poetry_dir
+    deactivate
+  fi
+
+  # Activate the environment if in a Poetry directory and no environment is currently active
+  if [[ $in_poetry_dir -eq 1 ]] && [[ $poetry_active -ne 1 ]]; then
+    venv_dir=$(poetry env info --path 2>/dev/null)
+    if [[ -n "$venv_dir" ]]; then
+      export poetry_active=1
+      export poetry_dir="$PWD"
+      source "${venv_dir}/bin/activate"
     fi
   fi
 }
 autoload -U add-zsh-hook
 add-zsh-hook chpwd _togglePoetryShell
-_togglePoetryShell
+_togglePoetryShell # Initial call to check the current directory at shell startup