Export VIRTUAL_ENV_PROMPT in activation scripts (#2194) (#2606)

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`.

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:
 

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
 

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%

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

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

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

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
       }
     })
 

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)
     }
 }

src/virtualenv/activation/python/activate_this.py

@@ -23,6 +23,7 @@
 # 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)

tests/unit/activation/conftest.py

@@ -114,36 +114,42 @@ def _get_test_lines(self, activate_script):
         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
         ]
 
     def assert_output(self, out, raw, tmp_path):
         # 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):

tests/unit/activation/test_nushell.py

@@ -19,7 +19,7 @@ def __init__(self, session) -> None:
             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

tests/unit/activation/test_python_activator.py

@@ -41,6 +41,7 @@ def print_r(value):
                 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 print_r(value):
             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,24 +64,27 @@ def print_r(value):
         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) :]
             assert new_lib_paths == set(new_sys_path[: len(new_lib_paths)])
 
             # 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):