Virtual env disable prompt

src/tempyenv/cli.py

@@ -11,6 +11,8 @@
 
 
 class TemporaryVenvCreator:
+    BASH_COMPATIBLE_SHELLS = {"bash", "sh", "zsh", "ksh"}
+
     def __init__(self, python_exec=None):
         self.temp_dir = None
         self.venv_path = None
@@ -35,11 +37,30 @@ def create_virtual_environment(self):
     def load_virtual_environment(self):
         try:
             print(f"Virtual environment loading from {self.venv_path}")
-            #current_shell = os.environ.get("SHELL") # Currently just support bash
-            current_shell = "bash"
+            user_shell_name = os.path.basename(os.environ.get("SHELL", ""))
+            if user_shell_name in self.BASH_COMPATIBLE_SHELLS:
+                current_shell = user_shell_name
+            else:
+                if shutil.which("bash") is None:
+                    raise RuntimeError(
+                        f"tempyenv only supports bash-compatible shells, but your shell "
+                        f"({user_shell_name or 'unknown'}) isn't one and bash was not found on PATH."
+                    )
+                print("Warning: tempyenv only supports bash-compatible shells, will attempt to launch a new bash shell")
+                current_shell = "bash"
+            current_ps1 = os.environ.get("PS1", "")
+            if os.environ.get("VIRTUAL_ENV_DISABLE_PROMPT"):
+                # User already opted out of venv prompt decoration - honor it.
+                set_ps1 = ""
+            elif current_ps1.startswith("(tempyenv)"):
+                # Already inside a tempyenv shell (e.g. tempyenv was run again
+                # from a nested shell) - avoid stacking duplicate prefixes.
+                set_ps1 = ""
+            else:
+                set_ps1 = "export PS1=\"(tempyenv)$PS1\\$ \" && "
             subprocess.run([f"{current_shell}",
                             "-c",
-                            f"export VIRTUAL_ENV_DISABLE_PROMPT=1 && source {self.venv_path}/bin/activate && export PS1=\"(tempyenv)$PS1\\$ \" && {current_shell}"],
+                            f"export VIRTUAL_ENV_DISABLE_PROMPT=1 && source {self.venv_path}/bin/activate && {set_ps1}{current_shell}"],
                             stdin=sys.stdin,
                             stdout=sys.stdout,
                             stderr=sys.stderr,

tests/units/test_cli.py

@@ -1,6 +1,8 @@
 import subprocess
 import sys
 
+import pytest
+
 from tempyenv.cli import TemporaryVenvCreator
 
 
@@ -75,6 +77,78 @@ def test_falls_back_to_stdlib_venv_when_uv_unavailable(monkeypatch):
     assert calls == [(["python3.10", "-m", "venv", "/tmp/fake/venv"], {"check": True})]
 
 
+def test_load_virtual_environment_sets_ps1_when_not_already_set(monkeypatch):
+    monkeypatch.delenv("PS1", raising=False)
+    calls = []
+    monkeypatch.setattr(
+        "tempyenv.cli.subprocess.run",
+        lambda _cmd, **_kwargs: calls.append((_cmd, _kwargs)),
+    )
+    cli = TemporaryVenvCreator()
+    cli.venv_path = "/tmp/fake/venv"
+    cli.load_virtual_environment()
+    cmd = calls[0][0]
+    assert 'export PS1="(tempyenv)$PS1\\$ "' in cmd[2]
+
+
+def test_load_virtual_environment_skips_ps1_when_already_set(monkeypatch):
+    monkeypatch.setenv("PS1", "(tempyenv)$ ")
+    calls = []
+    monkeypatch.setattr(
+        "tempyenv.cli.subprocess.run",
+        lambda _cmd, **_kwargs: calls.append((_cmd, _kwargs)),
+    )
+    cli = TemporaryVenvCreator()
+    cli.venv_path = "/tmp/fake/venv"
+    cli.load_virtual_environment()
+    cmd = calls[0][0]
+    assert 'export PS1=' not in cmd[2]
+
+
+def test_load_virtual_environment_honors_virtual_env_disable_prompt(monkeypatch):
+    monkeypatch.delenv("PS1", raising=False)
+    monkeypatch.setenv("VIRTUAL_ENV_DISABLE_PROMPT", "1")
+    calls = []
+    monkeypatch.setattr(
+        "tempyenv.cli.subprocess.run",
+        lambda _cmd, **_kwargs: calls.append((_cmd, _kwargs)),
+    )
+    cli = TemporaryVenvCreator()
+    cli.venv_path = "/tmp/fake/venv"
+    cli.load_virtual_environment()
+    cmd = calls[0][0]
+    assert 'export PS1=' not in cmd[2]
+
+
+def test_load_virtual_environment_uses_user_shell_when_supported(monkeypatch):
+    monkeypatch.setenv("SHELL", "/usr/bin/zsh")
+    calls = []
+    monkeypatch.setattr(
+        "tempyenv.cli.subprocess.run",
+        lambda _cmd, **_kwargs: calls.append((_cmd, _kwargs)),
+    )
+    cli = TemporaryVenvCreator()
+    cli.venv_path = "/tmp/fake/venv"
+    cli.load_virtual_environment()
+    cmd = calls[0][0]
+    assert cmd[0] == "zsh"
+
+
+def test_load_virtual_environment_raises_when_shell_unsupported_and_bash_missing(monkeypatch):
+    monkeypatch.setenv("SHELL", "/usr/bin/fish")
+    monkeypatch.setattr("tempyenv.cli.shutil.which", lambda _name: None)
+    calls = []
+    monkeypatch.setattr(
+        "tempyenv.cli.subprocess.run",
+        lambda _cmd, **_kwargs: calls.append((_cmd, _kwargs)),
+    )
+    cli = TemporaryVenvCreator()
+    cli.venv_path = "/tmp/fake/venv"
+    with pytest.raises(RuntimeError, match="bash-compatible shells"):
+        cli.load_virtual_environment()
+    assert calls == []
+
+
 def test_logs_error_when_uv_venv_fails(monkeypatch, caplog):
     monkeypatch.setattr("tempyenv.cli.shutil.which", lambda _name: "/usr/local/bin/uv")