Repository: pypa/virtualenv
Pull request #2606 (2023-07-13 23:15:07)

Export `VIRTUAL_ENV_PROMPT` in activation scripts (#2194)
See issue #2194. This change exports a `VIRTUAL_ENV_PROMPT` environment variable during venv activation so that users can use this in their own custom shell prompts if they like. ~I didn't add a test for this since I don't see any similar tests, and I'm not quite sure the bet way to go about writing one...~

I manually tested these changes in Bash and the results look good to me, though my changes for the other shell scripts are speculative. I did my best to guess how they should work, but there may be bugs!

---

### Thanks for contributing, make sure you address all the checklists (for details on how see [development documentation](https://virtualenv.pypa.io/en/latest/development.html#development))

- [x] ran the linter to address style issues (`tox -e fix`)
- [x] wrote descriptive pull request text
- [x] ensured there are test(s) validating the fix
- [x] added news fragment in `docs/changelog` folder
- [x] updated/extended the documentation


jimporter: > Please add extend the test suite to test for this.

I'll take a look. I'm going to mark this as a draft for now until I get tests resolved (and ideally for my second `venv` PR to merge so that I can be sure the behavior is the same between `venv` and `virtualenv`).
jimporter: @gaborbernat Ok, hopefully these tests pass. I had to make some educated guesses for a few of the shells.
jimporter: @gaborbernat Ok, I think these tests should be good now. One small note is that I renamed `VIRTUAL_PROMPT` to `VIRTUAL_PREFIX` in Nushell to reduce ambiguity, though I could revert that. As far as I can tell, that variable is only there to make it easier to write a unit test (this is the variable that holds `(<env-name>) `, as opposed to `VIRTUAL_ENV_PROMPT` which holds `<env-name>`).

Patch:
diff --git a/docs/changelog/2194.feature.rst b/docs/changelog/2194.feature.rst
new file mode 100644
index 00000000..e51ebecd
--- /dev/null
+++ b/docs/changelog/2194.feature.rst
@@ -0,0 +1,2 @@
+Export the prompt prefix as ``VIRTUAL_ENV_PROMPT`` when activating a virtual
+environment - by :user:`jimporter`.
diff --git a/docs/user_guide.rst b/docs/user_guide.rst
index fe64ed84..2169dfe0 100644
--- a/docs/user_guide.rst
+++ b/docs/user_guide.rst
@@ -242,8 +242,10 @@ Note that you don't have to activate a virtual environment to use it. You can in
 executables, rather than relying on your shell to resolve them to your virtual environment.
 
 Activator scripts also modify your shell prompt to indicate which environment is currently active, by prepending the
-environment name in brackets, like ``(venv)``. You can disable this behaviour by setting the environment variable
-``VIRTUAL_ENV_DISABLE_PROMPT`` to any value.
+environment name (or the name specified by ``--prompt`` when initially creating the environment) in brackets, like
+``(venv)``. You can disable this behaviour by setting the environment variable ``VIRTUAL_ENV_DISABLE_PROMPT`` to any
+value. You can also get the environment name via the environment variable ``VIRTUAL_ENV_PROMPT`` if you want to
+customize your prompt, for example.
 
 The scripts also provision a ``deactivate`` command that will allow you to undo the operation:
 
diff --git a/src/virtualenv/activation/bash/activate.sh b/src/virtualenv/activation/bash/activate.sh
index fb40db63..b06e3fd3 100644
--- a/src/virtualenv/activation/bash/activate.sh
+++ b/src/virtualenv/activation/bash/activate.sh
@@ -35,6 +35,7 @@ deactivate () {
     fi
 
     unset VIRTUAL_ENV
+    unset VIRTUAL_ENV_PROMPT
     if [ ! "${1-}" = "nondestructive" ] ; then
     # Self destruct!
         unset -f deactivate
@@ -54,6 +55,13 @@ _OLD_VIRTUAL_PATH="$PATH"
 PATH="$VIRTUAL_ENV/__BIN_NAME__:$PATH"
 export PATH
 
+if [ "x__VIRTUAL_PROMPT__" != x ] ; then
+    VIRTUAL_ENV_PROMPT="__VIRTUAL_PROMPT__"
+else
+    VIRTUAL_ENV_PROMPT=$(basename "$VIRTUAL_ENV")
+fi
+export VIRTUAL_ENV_PROMPT
+
 # unset PYTHONHOME if set
 if ! [ -z "${PYTHONHOME+_}" ] ; then
     _OLD_VIRTUAL_PYTHONHOME="$PYTHONHOME"
@@ -62,11 +70,7 @@ fi
 
 if [ -z "${VIRTUAL_ENV_DISABLE_PROMPT-}" ] ; then
     _OLD_VIRTUAL_PS1="${PS1-}"
-    if [ "x__VIRTUAL_PROMPT__" != x ] ; then
-        PS1="(__VIRTUAL_PROMPT__) ${PS1-}"
-    else
-        PS1="(`basename \"$VIRTUAL_ENV\"`) ${PS1-}"
-    fi
+    PS1="(${VIRTUAL_ENV_PROMPT}) ${PS1-}"
     export PS1
 fi
 
diff --git a/src/virtualenv/activation/batch/activate.bat b/src/virtualenv/activation/batch/activate.bat
index 816853c6..0bad4c1e 100644
--- a/src/virtualenv/activation/batch/activate.bat
+++ b/src/virtualenv/activation/batch/activate.bat
@@ -1,5 +1,10 @@
 @set "VIRTUAL_ENV=__VIRTUAL_ENV__"
 
+@set "VIRTUAL_ENV_PROMPT=__VIRTUAL_PROMPT__"
+@if NOT DEFINED VIRTUAL_ENV_PROMPT (
+    @for %%d in ("%VIRTUAL_ENV%") do @set "VIRTUAL_ENV_PROMPT=%%~nxd"
+)
+
 @if defined _OLD_VIRTUAL_PROMPT (
     @set "PROMPT=%_OLD_VIRTUAL_PROMPT%"
 ) else (
@@ -11,11 +16,7 @@
     )
 )
 @if not defined VIRTUAL_ENV_DISABLE_PROMPT (
-    @if "__VIRTUAL_PROMPT__" NEQ "" (
-        @set "PROMPT=(__VIRTUAL_PROMPT__) %PROMPT%"
-    ) else (
-        @for %%d in ("%VIRTUAL_ENV%") do @set "PROMPT=(%%~nxd) %PROMPT%"
-    )
+    @set "PROMPT=(%VIRTUAL_ENV_PROMPT%) %PROMPT%"
 )
 
 @REM Don't use () to avoid problems with them in %PATH%
diff --git a/src/virtualenv/activation/batch/deactivate.bat b/src/virtualenv/activation/batch/deactivate.bat
index f800f42c..8939c6c0 100644
--- a/src/virtualenv/activation/batch/deactivate.bat
+++ b/src/virtualenv/activation/batch/deactivate.bat
@@ -1,4 +1,5 @@
 @set VIRTUAL_ENV=
+@set VIRTUAL_ENV_PROMPT=
 
 @REM Don't use () to avoid problems with them in %PATH%
 @if not defined _OLD_VIRTUAL_PROMPT @goto ENDIFVPROMPT
diff --git a/src/virtualenv/activation/cshell/activate.csh b/src/virtualenv/activation/cshell/activate.csh
index 837dcda8..f0c9cca9 100644
--- a/src/virtualenv/activation/cshell/activate.csh
+++ b/src/virtualenv/activation/cshell/activate.csh
@@ -5,7 +5,7 @@
 set newline='\
 '
 
-alias deactivate 'test $?_OLD_VIRTUAL_PATH != 0 && setenv PATH "$_OLD_VIRTUAL_PATH:q" && unset _OLD_VIRTUAL_PATH; rehash; test $?_OLD_VIRTUAL_PROMPT != 0 && set prompt="$_OLD_VIRTUAL_PROMPT:q" && unset _OLD_VIRTUAL_PROMPT; unsetenv VIRTUAL_ENV; test "\!:*" != "nondestructive" && unalias deactivate && unalias pydoc'
+alias deactivate 'test $?_OLD_VIRTUAL_PATH != 0 && setenv PATH "$_OLD_VIRTUAL_PATH:q" && unset _OLD_VIRTUAL_PATH; rehash; test $?_OLD_VIRTUAL_PROMPT != 0 && set prompt="$_OLD_VIRTUAL_PROMPT:q" && unset _OLD_VIRTUAL_PROMPT; unsetenv VIRTUAL_ENV; unsetenv VIRTUAL_ENV_PROMPT; test "\!:*" != "nondestructive" && unalias deactivate && unalias pydoc'
 
 # Unset irrelevant variables.
 deactivate nondestructive
@@ -18,9 +18,9 @@ setenv PATH "$VIRTUAL_ENV:q/__BIN_NAME__:$PATH:q"
 
 
 if ('__VIRTUAL_PROMPT__' != "") then
-    set env_name = '(__VIRTUAL_PROMPT__) '
+    setenv VIRTUAL_ENV_PROMPT '__VIRTUAL_PROMPT__'
 else
-    set env_name = '('"$VIRTUAL_ENV:t:q"') '
+    setenv VIRTUAL_ENV_PROMPT "$VIRTUAL_ENV:t:q"
 endif
 
 if ( $?VIRTUAL_ENV_DISABLE_PROMPT ) then
@@ -42,7 +42,7 @@ if ( $do_prompt == "1" ) then
         if ( "$prompt:q" =~ *"$newline:q"* ) then
             :
         else
-            set prompt = "$env_name:q$prompt:q"
+            set prompt = '('"$VIRTUAL_ENV_PROMPT:q"') '"$prompt:q"
         endif
     endif
 endif
diff --git a/src/virtualenv/activation/fish/activate.fish b/src/virtualenv/activation/fish/activate.fish
index 62f631ea..fcedde46 100644
--- a/src/virtualenv/activation/fish/activate.fish
+++ b/src/virtualenv/activation/fish/activate.fish
@@ -44,6 +44,7 @@ function deactivate -d 'Exit virtualenv mode and return to the normal environmen
     end
 
     set -e VIRTUAL_ENV
+    set -e VIRTUAL_ENV_PROMPT
 
     if test "$argv[1]" != 'nondestructive'
         # Self-destruct!
@@ -67,6 +68,14 @@ else
 end
 set -gx PATH "$VIRTUAL_ENV"'/__BIN_NAME__' $PATH
 
+# Prompt override provided?
+# If not, just use the environment name.
+if test -n '__VIRTUAL_PROMPT__'
+    set -gx VIRTUAL_ENV_PROMPT '__VIRTUAL_PROMPT__'
+else
+    set -gx VIRTUAL_ENV_PROMPT (basename "$VIRTUAL_ENV")
+end
+
 # Unset `$PYTHONHOME` if set.
 if set -q PYTHONHOME
     set -gx _OLD_VIRTUAL_PYTHONHOME $PYTHONHOME
@@ -85,13 +94,7 @@ if test -z "$VIRTUAL_ENV_DISABLE_PROMPT"
         # Run the user's prompt first; it might depend on (pipe)status.
         set -l prompt (_old_fish_prompt)
 
-        # Prompt override provided?
-        # If not, just prepend the environment name.
-        if test -n '__VIRTUAL_PROMPT__'
-            printf '(%s) ' '__VIRTUAL_PROMPT__'
-        else
-            printf '(%s) ' (basename "$VIRTUAL_ENV")
-        end
+        printf '(%s) ' $VIRTUAL_ENV_PROMPT
 
         string join -- \n $prompt # handle multi-line prompts
     end
diff --git a/src/virtualenv/activation/nushell/activate.nu b/src/virtualenv/activation/nushell/activate.nu
index d70cbab4..19d4fa1d 100644
--- a/src/virtualenv/activation/nushell/activate.nu
+++ b/src/virtualenv/activation/nushell/activate.nu
@@ -46,20 +46,24 @@ export-env {
     let venv_path = ([$virtual_env $bin] | path join)
     let new_path = ($env | get $path_name | prepend $venv_path)
 
+    # If there is no default prompt, then use the env name instead
+    let virtual_env_prompt = (if ('__VIRTUAL_PROMPT__' | is-empty) {
+        ($virtual_env | path basename)
+    } else {
+        '__VIRTUAL_PROMPT__'
+    })
+
     let new_env = {
-        $path_name  : $new_path
-        VIRTUAL_ENV : $virtual_env
+        $path_name         : $new_path
+        VIRTUAL_ENV        : $virtual_env
+        VIRTUAL_ENV_PROMPT : $virtual_env_prompt
     }
 
     let new_env = (if (is-env-true 'VIRTUAL_ENV_DISABLE_PROMPT') {
       $new_env
     } else {
       # Creating the new prompt for the session
-      let virtual_prompt = (if ('__VIRTUAL_PROMPT__' | is-empty) {
-          $'(char lparen)($virtual_env | path basename)(char rparen) '
-      } else {
-          '(__VIRTUAL_PROMPT__) '
-      })
+      let virtual_prefix = $'(char lparen)($virtual_env_prompt)(char rparen) '
 
       # Back up the old prompt builder
       let old_prompt_command = (if (has-env 'PROMPT_COMMAND') {
@@ -68,20 +72,19 @@ export-env {
               ''
         })
 
-      # If there is no default prompt, then only the env is printed in the prompt
       let new_prompt = (if (has-env 'PROMPT_COMMAND') {
           if 'closure' in ($old_prompt_command | describe) {
-              {|| $'($virtual_prompt)(do $old_prompt_command)' }
+              {|| $'($virtual_prefix)(do $old_prompt_command)' }
           } else {
-              {|| $'($virtual_prompt)($old_prompt_command)' }
+              {|| $'($virtual_prefix)($old_prompt_command)' }
           }
       } else {
-          {|| $'($virtual_prompt)' }
+          {|| $'($virtual_prefix)' }
       })
 
       $new_env | merge {
         PROMPT_COMMAND      : $new_prompt
-        VIRTUAL_PROMPT      : $virtual_prompt
+        VIRTUAL_PREFIX      : $virtual_prefix
       }
     })
 
diff --git a/src/virtualenv/activation/powershell/activate.ps1 b/src/virtualenv/activation/powershell/activate.ps1
index d5243475..5ccfe120 100644
--- a/src/virtualenv/activation/powershell/activate.ps1
+++ b/src/virtualenv/activation/powershell/activate.ps1
@@ -16,6 +16,10 @@ function global:deactivate([switch] $NonDestructive) {
         Remove-Item env:VIRTUAL_ENV -ErrorAction SilentlyContinue
     }
 
+    if ($env:VIRTUAL_ENV_PROMPT) {
+        Remove-Item env:VIRTUAL_ENV_PROMPT -ErrorAction SilentlyContinue
+    }
+
     if (!$NonDestructive) {
         # Self destruct!
         Remove-Item function:deactivate
@@ -33,6 +37,13 @@ deactivate -nondestructive
 $VIRTUAL_ENV = $BASE_DIR
 $env:VIRTUAL_ENV = $VIRTUAL_ENV
 
+if ("__VIRTUAL_PROMPT__" -ne "") {
+    $env:VIRTUAL_ENV_PROMPT = "__VIRTUAL_PROMPT__"
+}
+else {
+    $env:VIRTUAL_ENV_PROMPT = $( Split-Path $env:VIRTUAL_ENV -Leaf )
+}
+
 New-Variable -Scope global -Name _OLD_VIRTUAL_PATH -Value $env:PATH
 
 $env:PATH = "$env:VIRTUAL_ENV/__BIN_NAME____PATH_SEP__" + $env:PATH
@@ -42,19 +53,9 @@ if (!$env:VIRTUAL_ENV_DISABLE_PROMPT) {
     }
     $function:_old_virtual_prompt = $function:prompt
 
-    if ("__VIRTUAL_PROMPT__" -ne "") {
-        function global:prompt {
-            # Add the custom prefix to the existing prompt
-            $previous_prompt_value = & $function:_old_virtual_prompt
-            ("(__VIRTUAL_PROMPT__) " + $previous_prompt_value)
-        }
-    }
-    else {
-        function global:prompt {
-            # Add a prefix to the current prompt, but don't discard it.
-            $previous_prompt_value = & $function:_old_virtual_prompt
-            $new_prompt_value = "($( Split-Path $env:VIRTUAL_ENV -Leaf )) "
-            ($new_prompt_value + $previous_prompt_value)
-        }
+    function global:prompt {
+        # Add the custom prefix to the existing prompt
+        $previous_prompt_value = & $function:_old_virtual_prompt
+        ("(" + $env:VIRTUAL_ENV_PROMPT + ") " + $previous_prompt_value)
     }
 }
diff --git a/src/virtualenv/activation/python/activate_this.py b/src/virtualenv/activation/python/activate_this.py
index d01bd58e..33a3eaee 100644
--- a/src/virtualenv/activation/python/activate_this.py
+++ b/src/virtualenv/activation/python/activate_this.py
@@ -23,6 +23,7 @@ base = bin_dir[: -len("__BIN_NAME__") - 1]  # strip away the bin part from the _
 # prepend bin to PATH (this file is inside the bin directory)
 os.environ["PATH"] = os.pathsep.join([bin_dir, *os.environ.get("PATH", "").split(os.pathsep)])
 os.environ["VIRTUAL_ENV"] = base  # virtual env is right above bin directory
+os.environ["VIRTUAL_ENV_PROMPT"] = "__VIRTUAL_PROMPT__" or os.path.basename(base)  # noqa: SIM222
 
 # add the virtual environments libraries to the host python import mechanism
 prev_length = len(sys.path)


Test patch:
diff --git a/tests/unit/activation/conftest.py b/tests/unit/activation/conftest.py
index 64c481da..04946f9c 100644
--- a/tests/unit/activation/conftest.py
+++ b/tests/unit/activation/conftest.py
@@ -114,15 +114,18 @@ class ActivationTester:
         return [
             self.print_python_exe(),
             self.print_os_env_var("VIRTUAL_ENV"),
+            self.print_os_env_var("VIRTUAL_ENV_PROMPT"),
             self.activate_call(activate_script),
             self.print_python_exe(),
             self.print_os_env_var("VIRTUAL_ENV"),
+            self.print_os_env_var("VIRTUAL_ENV_PROMPT"),
             self.print_prompt(),
             # \\ loads documentation from the virtualenv site packages
             self.pydoc_call,
             self.deactivate,
             self.print_python_exe(),
             self.print_os_env_var("VIRTUAL_ENV"),
+            self.print_os_env_var("VIRTUAL_ENV_PROMPT"),
             "",  # just finish with an empty new line
         ]
 
@@ -130,20 +133,23 @@ class ActivationTester:
         # pre-activation
         assert out[0], raw
         assert out[1] == "None", raw
+        assert out[2] == "None", raw
         # post-activation
         expected = self._creator.exe.parent / os.path.basename(sys.executable)
-        assert self.norm_path(out[2]) == self.norm_path(expected), raw
-        assert self.norm_path(out[3]) == self.norm_path(self._creator.dest).replace("\\\\", "\\"), raw
+        assert self.norm_path(out[3]) == self.norm_path(expected), raw
+        assert self.norm_path(out[4]) == self.norm_path(self._creator.dest).replace("\\\\", "\\"), raw
+        assert out[5] == self._creator.env_name
         # Some attempts to test the prompt output print more than 1 line.
         # So we need to check if the prompt exists on any of them.
         prompt_text = f"({self._creator.env_name}) "
-        assert any(prompt_text in line for line in out[4:-3]), raw
+        assert any(prompt_text in line for line in out[6:-4]), raw
 
-        assert out[-3] == "wrote pydoc_test.html", raw
+        assert out[-4] == "wrote pydoc_test.html", raw
         content = tmp_path / "pydoc_test.html"
         assert content.exists(), raw
         # post deactivation, same as before
-        assert out[-2] == out[0], raw
+        assert out[-3] == out[0], raw
+        assert out[-2] == "None", raw
         assert out[-1] == "None", raw
 
     def quote(self, s):
diff --git a/tests/unit/activation/test_nushell.py b/tests/unit/activation/test_nushell.py
index ed2ec79d..fbf75e39 100644
--- a/tests/unit/activation/test_nushell.py
+++ b/tests/unit/activation/test_nushell.py
@@ -19,7 +19,7 @@ def test_nushell(activation_tester_class, activation_tester):
             self.unix_line_ending = not IS_WIN
 
         def print_prompt(self):
-            return r"print $env.VIRTUAL_PROMPT"
+            return r"print $env.VIRTUAL_PREFIX"
 
         def activate_call(self, script):
             # Commands are called without quotes in Nushell
diff --git a/tests/unit/activation/test_python_activator.py b/tests/unit/activation/test_python_activator.py
index 17cda681..f773f74a 100644
--- a/tests/unit/activation/test_python_activator.py
+++ b/tests/unit/activation/test_python_activator.py
@@ -41,6 +41,7 @@ def test_python(raise_on_non_source_class, activation_tester):
                 print(repr(value))
 
             print_r(os.environ.get("VIRTUAL_ENV"))
+            print_r(os.environ.get("VIRTUAL_ENV_PROMPT"))
             print_r(os.environ.get("PATH").split(os.pathsep))
             print_r(sys.path)
 
@@ -51,6 +52,7 @@ def test_python(raise_on_non_source_class, activation_tester):
             exec(content, {{"__file__": file_at}})
 
             print_r(os.environ.get("VIRTUAL_ENV"))
+            print_r(os.environ.get("VIRTUAL_ENV_PROMPT"))
             print_r(os.environ.get("PATH").split(os.pathsep))
             print_r(sys.path)
 
@@ -62,16 +64,19 @@ def test_python(raise_on_non_source_class, activation_tester):
         def assert_output(self, out, raw, tmp_path):  # noqa: ARG002
             out = [literal_eval(i) for i in out]
             assert out[0] is None  # start with VIRTUAL_ENV None
+            assert out[1] is None  # likewise for VIRTUAL_ENV_PROMPT
 
-            prev_path = out[1]
-            prev_sys_path = out[2]
-            assert out[3] == str(self._creator.dest)  # VIRTUAL_ENV now points to the virtual env folder
+            prev_path = out[2]
+            prev_sys_path = out[3]
+            assert out[4] == str(self._creator.dest)  # VIRTUAL_ENV now points to the virtual env folder
 
-            new_path = out[4]  # PATH now starts with bin path of current
+            assert out[5] == str(self._creator.env_name)  # VIRTUAL_ENV_PROMPT now has the env name
+
+            new_path = out[6]  # PATH now starts with bin path of current
             assert ([str(self._creator.bin_dir), *prev_path]) == new_path
 
             # sys path contains the site package at its start
-            new_sys_path = out[5]
+            new_sys_path = out[7]
 
             new_lib_paths = {str(i) for i in self._creator.libs}
             assert prev_sys_path == new_sys_path[len(new_lib_paths) :]
@@ -79,7 +84,7 @@ def test_python(raise_on_non_source_class, activation_tester):
 
             # manage to import from activate site package
             dest = self.norm_path(self._creator.purelib / "pydoc_test.py")
-            found = self.norm_path(out[6])
+            found = self.norm_path(out[8])
             assert found.startswith(dest)
 
         def non_source_activate(self, activate_script):