From d5efae87cba2013603b9570612aa867ec39a2315 Mon Sep 17 00:00:00 2001 From: David Whiteside Date: Sun, 2 Aug 2026 09:39:03 -0600 Subject: [PATCH] add support for VIRTUAL_ENV_DISABLE_PROMPT to not change the prompt and better error handling if the user shell isnt bash --- src/tempyenv/cli.py | 27 +++++++++++++-- tests/units/test_cli.py | 74 +++++++++++++++++++++++++++++++++++++++++ 2 files changed, 98 insertions(+), 3 deletions(-) diff --git a/src/tempyenv/cli.py b/src/tempyenv/cli.py index b78f176..1c1a4f1 100644 --- a/src/tempyenv/cli.py +++ b/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, diff --git a/tests/units/test_cli.py b/tests/units/test_cli.py index b71fb26..bda058e 100644 --- a/tests/units/test_cli.py +++ b/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")