From d4c9655e26a10e03c664410a6653cf5073bf3bcb Mon Sep 17 00:00:00 2001
From: Jovier Jimenez <51138002+jimenezj8@users.noreply.github.com>
Date: Sun, 15 Sep 2024 02:40:44 -0400
Subject: [PATCH] Add prompt customization (#476)

---
 README.md                |  1 +
 bin/pyenv-sh-activate    |  7 ++++++-
 test/activate.bats       | 26 ++++++++++++++++++++++++++
 test/conda-activate.bats | 30 ++++++++++++++++++++++++++++++
 test/hooks.bats          |  1 +
 5 files changed, 64 insertions(+), 1 deletion(-)

diff --git a/README.md b/README.md
index d533888..01c18d0 100644
--- a/README.md
+++ b/README.md
@@ -243,6 +243,7 @@ You can set certain environment variables to control pyenv-virtualenv.
 * `PIP_VERSION`, if set and `venv` is preferred
   over `virtualenv`, install the specified version of pip.
 * `PYENV_VIRTUALENV_VERBOSE_ACTIVATE`, if set, shows some verbose outputs on activation and deactivation
+* `PYENV_VIRTUALENV_PROMPT`, if set, allows users to customize how `pyenv-virtualenv` modifies their shell prompt. The default prompt ("(venv)") is overwritten with any user-specified text. Specify the location of the virtual environment name with the string `{venv}`. For example, the default prompt string would be `({venv})`.
 
 ## Version History
 
diff --git a/bin/pyenv-sh-activate b/bin/pyenv-sh-activate
index 38091ee..5c53b50 100755
--- a/bin/pyenv-sh-activate
+++ b/bin/pyenv-sh-activate
@@ -262,9 +262,14 @@ EOS
     fi
     ;;
   * )
+    if [ -z "${PYENV_VIRTUALENV_PROMPT}" ]; then
+      PYENV_VIRTUALENV_PROMPT="(${venv})"
+    else
+      PYENV_VIRTUALENV_PROMPT="${PYENV_VIRTUALENV_PROMPT/\{venv\}/${venv}}"
+    fi
     cat <<EOS
 export _OLD_VIRTUAL_PS1="\${PS1:-}";
-export PS1="(${venv}) \${PS1:-}";
+export PS1="${PYENV_VIRTUALENV_PROMPT} \${PS1:-}";
 EOS
     ;;
   esac
diff --git a/test/activate.bats b/test/activate.bats
index 8c1760a..25f0f05 100644
--- a/test/activate.bats
+++ b/test/activate.bats
@@ -15,6 +15,7 @@ setup() {
   unset PYENV_VIRTUALENV_DISABLE_PROMPT
   unset PYENV_VIRTUAL_ENV_DISABLE_PROMPT
   unset VIRTUAL_ENV_DISABLE_PROMPT
+  unset PYENV_VIRTUALENV_PROMPT
   unset _OLD_VIRTUAL_PS1
   stub pyenv-hooks "activate : echo"
 }
@@ -44,6 +45,31 @@ EOS
   unstub pyenv-sh-deactivate
 }
 
+@test "activate virtualenv from current version with custom prompt" {
+  export PYENV_VIRTUALENV_INIT=1
+
+  stub pyenv-version-name "echo venv"
+  stub pyenv-virtualenv-prefix "venv : echo \"${PYENV_ROOT}/versions/venv\""
+  stub pyenv-prefix "venv : echo \"${PYENV_ROOT}/versions/venv\""
+  stub pyenv-sh-deactivate "--force --quiet : echo deactivated"
+
+  PYENV_SHELL="bash" PYENV_VERSION="venv" PYENV_VIRTUALENV_PROMPT='venv:{venv}' run pyenv-sh-activate
+
+  assert_success
+  assert_output <<EOS
+deactivated
+export PYENV_VIRTUAL_ENV="${PYENV_ROOT}/versions/venv";
+export VIRTUAL_ENV="${PYENV_ROOT}/versions/venv";
+export _OLD_VIRTUAL_PS1="\${PS1:-}";
+export PS1="venv:venv \${PS1:-}";
+EOS
+
+  unstub pyenv-version-name
+  unstub pyenv-virtualenv-prefix
+  unstub pyenv-prefix
+  unstub pyenv-sh-deactivate
+}
+
 @test "activate virtualenv from current version (quiet)" {
   export PYENV_VIRTUALENV_INIT=1
 
diff --git a/test/conda-activate.bats b/test/conda-activate.bats
index 452f6e3..51721c9 100644
--- a/test/conda-activate.bats
+++ b/test/conda-activate.bats
@@ -16,6 +16,7 @@ setup() {
   unset PYENV_VIRTUALENV_DISABLE_PROMPT
   unset PYENV_VIRTUAL_ENV_DISABLE_PROMPT
   unset VIRTUAL_ENV_DISABLE_PROMPT
+  unset PYENV_VIRTUALENV_PROMPT
   unset _OLD_VIRTUAL_PS1
   stub pyenv-hooks "activate : echo"
 }
@@ -53,6 +54,35 @@ EOS
   teardown_conda "anaconda-2.3.0"
 }
 
+@test "activate conda root from current version with custom prompt" {
+  export PYENV_VIRTUALENV_INIT=1
+
+  setup_conda "anaconda-2.3.0"
+  stub pyenv-version-name "echo anaconda-2.3.0"
+  stub pyenv-virtualenv-prefix "anaconda-2.3.0 : echo \"${PYENV_ROOT}/versions/anaconda-2.3.0\""
+  stub pyenv-prefix "anaconda-2.3.0 : echo \"${PYENV_ROOT}/versions/anaconda-2.3.0\""
+  stub pyenv-sh-deactivate "--force --quiet : echo deactivated"
+
+  PYENV_SHELL="bash" PYENV_VERSION="anaconda-2.3.0" PYENV_VIRTUALENV_PROMPT='venv:{venv}' run pyenv-sh-activate
+
+  assert_success
+  assert_output <<EOS
+deactivated
+export PYENV_VIRTUAL_ENV="${PYENV_ROOT}/versions/anaconda-2.3.0";
+export VIRTUAL_ENV="${PYENV_ROOT}/versions/anaconda-2.3.0";
+export CONDA_DEFAULT_ENV="root";
+export _OLD_VIRTUAL_PS1="\${PS1:-}";
+export PS1="venv:anaconda-2.3.0 \${PS1:-}";
+export CONDA_PREFIX="${TMP}/pyenv/versions/anaconda-2.3.0";
+EOS
+
+  unstub pyenv-version-name
+  unstub pyenv-virtualenv-prefix
+  unstub pyenv-prefix
+  unstub pyenv-sh-deactivate
+  teardown_conda "anaconda-2.3.0"
+}
+
 @test "activate conda root from current version (fish)" {
   export PYENV_VIRTUALENV_INIT=1
 
diff --git a/test/hooks.bats b/test/hooks.bats
index 4149c2e..6ff00d1 100644
--- a/test/hooks.bats
+++ b/test/hooks.bats
@@ -6,6 +6,7 @@ setup() {
   export PYENV_ROOT="${TMP}/pyenv"
   export HOOK_PATH="${TMP}/i has hooks"
   mkdir -p "$HOOK_PATH"
+  unset PYENV_VIRTUALENV_PROMPT
 }
 
 @test "pyenv-virtualenv hooks" {