From d8d54c3863ef2737c655cf43052adc61bcfbe9f4 Mon Sep 17 00:00:00 2001 From: jaimergp Date: Wed, 15 Jan 2025 22:50:26 +0100 Subject: [PATCH] Add '--hook' option (#2) --- .devcontainer/post_create.sh | 1 - .github/workflows/test.yml | 3 +- README.md | 6 +- conda_spawn/__init__.py | 3 +- conda_spawn/cli.py | 45 +++++++- conda_spawn/exceptions.py | 4 +- conda_spawn/main.py | 12 ++- conda_spawn/plugin.py | 2 +- conda_spawn/shell.py | 204 +++++++++++++++++++++-------------- tests/conftest.py | 2 +- tests/test_cli.py | 1 + tests/test_shell.py | 90 +++++++++++++--- 12 files changed, 256 insertions(+), 117 deletions(-) diff --git a/.devcontainer/post_create.sh b/.devcontainer/post_create.sh index 77b1e42..c4e2fe9 100644 --- a/.devcontainer/post_create.sh +++ b/.devcontainer/post_create.sh @@ -28,4 +28,3 @@ if [ ! -f "$SRC_CONDA_SPAWN/pyproject.toml" ]; then echo "https://github.com/conda-incubator/conda-spawn not found! Please clone or mount to $SRC_CONDA_SPAWN" exit 1 fi - diff --git a/.github/workflows/test.yml b/.github/workflows/test.yml index 9130037..e423de8 100644 --- a/.github/workflows/test.yml +++ b/.github/workflows/test.yml @@ -51,7 +51,7 @@ jobs: pixi run --environment ${{ env.PIXI_ENV_NAME }} conda info - name: Run tests run: pixi run --environment ${{ env.PIXI_ENV_NAME }} test --basetemp=${{ runner.os == 'Windows' && 'D:\\temp' || runner.temp }} - + build-conda: name: Build conda package (${{ matrix.os }}) runs-on: ${{ matrix.os }} @@ -73,4 +73,3 @@ jobs: pixi run --environment build dev - name: Build recipe run: pixi run --environment build build - \ No newline at end of file diff --git a/README.md b/README.md index 7c93ee9..e4c0544 100644 --- a/README.md +++ b/README.md @@ -52,13 +52,13 @@ For in-script usage, please consider these replacements for `conda activate`: For Unix shell scripts: ```bash -eval "$(conda shell.posix activate )" +eval "$(conda spawn --hook --shell posix -n )" ``` For Windows CMD scripts: ```cmd -FOR /F "tokens=*" %%g IN ('conda shell.cmd.exe activate ') do @CALL %%g +FOR /F "tokens=*" %%g IN ('conda spawn --hook --shell cmd -n ') do @CALL %%g ``` For Windows Powershell scripts: @@ -72,7 +72,7 @@ For example, if you want to create a new environment and activate it, it would l ```bash # Assumes `conda` is in PATH conda create -n new-env python numpy -eval "$(conda shell.posix activate new-env)" +eval "$(conda spawn --hook --shell powershell -n new-env)" python -c "import numpy" ``` diff --git a/conda_spawn/__init__.py b/conda_spawn/__init__.py index 7317078..9065035 100644 --- a/conda_spawn/__init__.py +++ b/conda_spawn/__init__.py @@ -1,4 +1,5 @@ """ conda go: activate conda environments in new shell processes. """ -from .main import spawn + +from .main import spawn, hook # noqa diff --git a/conda_spawn/cli.py b/conda_spawn/cli.py index 9002843..ea9fc59 100644 --- a/conda_spawn/cli.py +++ b/conda_spawn/cli.py @@ -5,7 +5,9 @@ from __future__ import annotations import argparse +from textwrap import dedent +from conda.exceptions import ArgumentError from conda.cli.conda_argparse import ( add_parser_help, add_parser_prefix, @@ -13,26 +15,59 @@ def configure_parser(parser: argparse.ArgumentParser): + from .shell import SHELLS + add_parser_help(parser) add_parser_prefix(parser, prefix_required=True) parser.add_argument( "command", + metavar="COMMAND [args]", nargs="*", - help="Optional program and arguments to run after starting the shell.", + help="Optional program to run after starting the shell. " + "Use -- before the program if providing arguments.", ) shell_group = parser.add_argument_group("Shell options") shell_group.add_argument( - "-s", + "--hook", + action="store_true", + help=( + "Print the shell activation logic so it can be sourced in-process. " + "This is meant to be used in scripts only." + ), + ) + shell_group.add_argument( "--shell", - help="Shell to use for the new session. " - "If not specified, autodetect shell in use.", + choices=SHELLS, + help="Shell to use for the new session. If not specified, autodetect shell in use.", ) + parser.prog = "conda spawn" + parser.epilog = dedent( + """ + Examples for --hook usage in different shells: + POSIX: + source "$(conda spawn --hook -n ENV-NAME)" + CMD: + FOR /F "tokens=*" %%g IN ('conda spawn --hook -n ENV-NAME') do @CALL %%g + Powershell: + conda spawn --hook -n ENV-NAME | Out-String | Invoke-Expression + """ + ).lstrip() + def execute(args: argparse.Namespace) -> int: - from .main import spawn, environment_speficier_to_path, shell_specifier_to_shell + from .main import ( + hook, + spawn, + environment_speficier_to_path, + shell_specifier_to_shell, + ) prefix = environment_speficier_to_path(args.name, args.prefix) shell = shell_specifier_to_shell(args.shell) + if args.hook: + if args.command: + raise ArgumentError("COMMAND cannot be provided with --hook.") + return hook(prefix, shell) return spawn(prefix, shell, command=args.command) diff --git a/conda_spawn/exceptions.py b/conda_spawn/exceptions.py index 0c6b07f..52b5f82 100644 --- a/conda_spawn/exceptions.py +++ b/conda_spawn/exceptions.py @@ -1,5 +1,4 @@ -""" -""" +""" """ from conda.base.constants import COMPATIBLE_SHELLS from conda.common.io import dashlist @@ -14,4 +13,3 @@ def __init__(self, name: str): f"{dashlist(COMPATIBLE_SHELLS)}" ) super().__init__(message) - diff --git a/conda_spawn/main.py b/conda_spawn/main.py index b96c710..f87d70e 100644 --- a/conda_spawn/main.py +++ b/conda_spawn/main.py @@ -18,7 +18,17 @@ def spawn( ) -> int: if shell_cls is None: shell_cls = detect_shell_class() - return shell_cls().spawn(prefix, command=command) + return shell_cls(prefix).spawn(command=command) + + +def hook(prefix: Path, shell_cls: Shell | None = None) -> int: + if shell_cls is None: + shell_cls = detect_shell_class() + script = shell_cls(prefix).script() + prompt = shell_cls(prefix).prompt() + print(script) + print(prompt) + return 0 def environment_speficier_to_path( diff --git a/conda_spawn/plugin.py b/conda_spawn/plugin.py index 69d5ef9..71a670d 100644 --- a/conda_spawn/plugin.py +++ b/conda_spawn/plugin.py @@ -9,7 +9,7 @@ def conda_subcommands(): yield plugins.CondaSubcommand( name="spawn", - summary="activate conda environments in new shell processes", + summary="Activate conda environments in new shell processes.", action=cli.execute, configure_parser=cli.configure_parser, ) diff --git a/conda_spawn/shell.py b/conda_spawn/shell.py index acffb64..4535cdb 100644 --- a/conda_spawn/shell.py +++ b/conda_spawn/shell.py @@ -26,6 +26,14 @@ class Shell: + Activator: activate._Activator + + def __init__(self, prefix: Path): + self.prefix = prefix + self._prefix_str = str(prefix) + self._activator = self.Activator(["activate", str(self.prefix)]) + self._files_to_remove = [] + def spawn(self, prefix: Path) -> int: """ Creates a new shell session with the conda environment at `path` @@ -35,12 +43,64 @@ def spawn(self, prefix: Path) -> int: """ raise NotImplementedError + def script(self) -> str: + raise NotImplementedError + + def prompt(self) -> str: + raise NotImplementedError + + def prompt_modifier(self) -> str: + conda_default_env = os.getenv( + "CONDA_DEFAULT_ENV", self._activator._default_env(self._prefix_str) + ) + return self._activator._prompt_modifier(self._prefix_str, conda_default_env) + + def executable(self) -> str: + raise NotImplementedError + + def args(self) -> tuple[str, ...]: + raise NotImplementedError + + def env(self) -> dict[str, str]: + env = os.environ.copy() + env["CONDA_SPAWN"] = "1" + return env + + def __del__(self): + for path in self._files_to_remove: + try: + os.unlink(path) + except OSError as exc: + log.debug("Could not delete %s", path, exc_info=exc) + class PosixShell(Shell): Activator = activate.PosixActivator - tmp_suffix = ".sh" + default_shell = "/bin/sh" + default_args = ("-l", "-i") - def spawn_tty(self, prefix: Path, command: Iterable[str] | None = None) -> pexpect.spawn: + def spawn(self, command: Iterable[str] | None = None) -> int: + return self.spawn_tty(command).wait() + + def script(self) -> str: + script = self._activator.execute() + lines = [] + for line in script.splitlines(keepends=True): + if "PS1=" in line: + continue + lines.append(line) + return "".join(lines) + + def prompt(self) -> str: + return f'PS1="{self.prompt_modifier()}${{PS1:-}}"' + + def executable(self): + return os.environ.get("SHELL", self.default_shell) + + def args(self): + return self.default_args + + def spawn_tty(self, command: Iterable[str] | None = None) -> pexpect.spawn: def _sigwinch_passthrough(sig, data): # NOTE: Taken verbatim from pexpect's .interact() docstring. # Check for buggy platforms (see pexpect.setwinsize()). @@ -52,13 +112,12 @@ def _sigwinch_passthrough(sig, data): a = struct.unpack("HHHH", fcntl.ioctl(sys.stdout.fileno(), TIOCGWINSZ, s)) child.setwinsize(a[0], a[1]) - script, prompt = self.script_and_prompt(prefix) - executable, args = self.executable_and_args() size = shutil.get_terminal_size() + executable = self.executable() child = pexpect.spawn( - executable, - args, + self.executable(), + [*self.args()], env=self.env(), echo=False, dimensions=(size.lines, size.columns), @@ -66,17 +125,21 @@ def _sigwinch_passthrough(sig, data): try: with NamedTemporaryFile( prefix="conda-spawn-", - suffix=self.tmp_suffix, + suffix=self.Activator.script_extension, delete=False, mode="w", ) as f: - f.write(script) + f.write(self.script()) signal.signal(signal.SIGWINCH, _sigwinch_passthrough) # Source the activation script. We do this in a single line for performance. - # It's slower to send several lines than paying the IO overhead. - child.sendline(f' . "{f.name}" && PS1="{prompt}${{PS1:-}}" && stty echo') + # (It's slower to send several lines than paying the IO overhead). + # We set the PS1 prompt outside the script because it's otherwise invisible. + # stty echo is equivalent to `child.setecho(True)` but the latter didn't work + # reliably across all shells and OSs. + child.sendline(f' . "{f.name}" && {self.prompt()} && stty echo') os.read(child.child_fd, 4096) # consume buffer before interact if Path(executable).name == "zsh": + # zsh also needs this for a truly silent activation child.expect("\r\n") if command: child.sendline(shlex.join(command)) @@ -84,131 +147,107 @@ def _sigwinch_passthrough(sig, data): child.interact() return child finally: - os.unlink(f.name) - - def spawn(self, prefix: Path, command: Iterable[str] | None = None) -> int: - return self.spawn_tty(prefix, command).wait() + self._files_to_remove.append(f.name) - def script_and_prompt(self, prefix: Path) -> tuple[str, str]: - activator = self.Activator(["activate", str(prefix)]) - conda_default_env = os.getenv( - "CONDA_DEFAULT_ENV", activator._default_env(str(prefix)) - ) - prompt = activator._prompt_modifier(str(prefix), conda_default_env) - script = activator.execute() - lines = [] - for line in script.splitlines(keepends=True): - if "PS1=" in line: - continue - lines.append(line) - script = "".join(lines) - return script, prompt - def executable_and_args(self) -> tuple[str, list[str]]: - # TODO: Customize which shell gets used; this below is the default! - return os.environ.get("SHELL", "/bin/bash"), ["-l", "-i"] +class BashShell(PosixShell): + def executable(self): + return "bash" - def env(self) -> dict[str, str]: - env = os.environ.copy() - env["CONDA_SPAWN"] = "1" - return env + +class ZshShell(PosixShell): + def executable(self): + return "zsh" class CshShell(Shell): - def spawn(self, prefix: Path, command: Iterable[str] | None = None) -> int: ... + pass class XonshShell(Shell): - def spawn(self, prefix: Path, command: Iterable[str] | None = None) -> int: ... + pass class FishShell(Shell): - def spawn(self, prefix: Path, command: Iterable[str] | None = None) -> int: ... + pass class PowershellShell(Shell): Activator = activate.PowerShellActivator - tmp_suffix = ".ps1" - def spawn_popen(self, prefix: Path, command: Iterable[str] | None = None, **kwargs) -> subprocess.Popen: - executable, args = self.executable_and_args() - script, _ = self.script_and_prompt(prefix) + def spawn_popen( + self, command: Iterable[str] | None = None, **kwargs + ) -> subprocess.Popen: try: with NamedTemporaryFile( prefix="conda-spawn-", - suffix=self.tmp_suffix, + suffix=self.Activator.script_extension, delete=False, mode="w", ) as f: - f.write(f"{script}\r\n") + f.write(f"{self.script()}\r\n") + f.write(f"{self.prompt()}\r\n") if command: command = subprocess.list2cmdline(command) f.write(f"echo {command}\r\n") f.write(f"{command}\r\n") - return subprocess.Popen([executable, *args, f.name], env=self.env(), **kwargs) + return subprocess.Popen( + [self.executable(), *self.args(), f.name], env=self.env(), **kwargs + ) finally: - self._tmpfile = f.name + self._files_to_remove.append(f.name) - def spawn(self, prefix: Path, command: Iterable[str] | None = None) -> int: - proc = self.spawn_popen(prefix, command) + def spawn(self, command: Iterable[str] | None = None) -> int: + proc = self.spawn_popen(command) proc.communicate() return proc.wait() - def script_and_prompt(self, prefix: Path) -> tuple[str, str]: - activator = self.Activator(["activate", str(prefix)]) - conda_default_env = os.getenv( - "CONDA_DEFAULT_ENV", activator._default_env(str(prefix)) - ) - prompt_mod = activator._prompt_modifier(str(prefix), conda_default_env) - script = activator.execute() - script += ( + def script(self) -> str: + return self._activator.execute() + + def prompt(self) -> str: + return ( "\r\n$old_prompt = $function:prompt\r\n" - f'function prompt {{"{prompt_mod}$($old_prompt.Invoke())"}};' + f'function prompt {{"{self.prompt_modifier()}$($old_prompt.Invoke())"}};' ) - return script, "" - def executable_and_args(self) -> tuple[str, list[str]]: - # TODO: Customize which shell gets used; this below is the default! - return "powershell", ["-NoLogo", "-NoExit", "-File"] + def executable(self) -> str: + return "powershell" + + def args(self) -> tuple[str, ...]: + return ("-NoLogo", "-NoExit", "-File") def env(self) -> dict[str, str]: env = os.environ.copy() env["CONDA_SPAWN"] = "1" return env - - def __del__(self): - if getattr(self, "_tmpfile", None): - os.unlink(self._tmpfile) class CmdExeShell(PowershellShell): Activator = activate.CmdExeActivator - tmp_suffix = ".bat" - def script_and_prompt(self, prefix: Path) -> tuple[str, str]: - activator = self.Activator(["activate", str(prefix)]) - conda_default_env = os.getenv( - "CONDA_DEFAULT_ENV", activator._default_env(str(prefix)) - ) - prompt_mod = activator._prompt_modifier(str(prefix), conda_default_env) - script = "\r\n".join( + def script(self): + return "\r\n".join( [ "@ECHO OFF", - Path(activator.execute()).read_text(), - f'@SET "PROMPT={prompt_mod}$P$G"', - "\r\n@ECHO ON\r\n" + Path(self._activator.execute()).read_text(), + "@ECHO ON", ] ) - return script, "" - def executable_and_args(self) -> tuple[str, list[str]]: - # TODO: Customize which shell gets used; this below is the default! - return "cmd", ["/K"] + def prompt(self) -> str: + return f'@SET "PROMPT={self.prompt_modifier()}$P$G"' + + def executable(self) -> str: + return "cmd" + + def args(self) -> tuple[str, ...]: + return ("/D", "/K") SHELLS: dict[str, type[Shell]] = { "ash": PosixShell, - "bash": PosixShell, + "bash": BashShell, "cmd.exe": CmdExeShell, "cmd": CmdExeShell, "csh": CshShell, @@ -216,9 +255,10 @@ def executable_and_args(self) -> tuple[str, list[str]]: "fish": FishShell, "posix": PosixShell, "powershell": PowershellShell, + "pwsh": PowershellShell, "tcsh": CshShell, "xonsh": XonshShell, - "zsh": PosixShell, + "zsh": ZshShell, } diff --git a/tests/conftest.py b/tests/conftest.py index 5801412..0d1d69c 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -1 +1 @@ -from conda.testing import conda_cli +pytest_plugins = ("conda.testing.fixtures",) diff --git a/tests/test_cli.py b/tests/test_cli.py index 6d433a8..205d823 100644 --- a/tests/test_cli.py +++ b/tests/test_cli.py @@ -1,5 +1,6 @@ import sys + def test_cli(monkeypatch, conda_cli): monkeypatch.setattr(sys, "argv", ["conda", *sys.argv[1:]]) out, err, _ = conda_cli("spawn", "-h", raises=SystemExit) diff --git a/tests/test_shell.py b/tests/test_shell.py index ba2efb7..6138280 100644 --- a/tests/test_shell.py +++ b/tests/test_shell.py @@ -1,38 +1,94 @@ -import signal import sys import pytest from conda_spawn.shell import PosixShell, PowershellShell, CmdExeShell -from subprocess import PIPE +from subprocess import PIPE, check_output + + +@pytest.fixture(scope="session") +def simple_env(session_tmp_env): + with session_tmp_env() as prefix: + yield prefix + @pytest.mark.skipif(sys.platform == "win32", reason="Pty's only available on Unix") -def test_posix_shell(): - proc = PosixShell().spawn_tty(sys.prefix) +def test_posix_shell(simple_env): + shell = PosixShell(simple_env) + proc = shell.spawn_tty() proc.sendline("env") - proc.expect("CONDA_SPAWN") - proc.sendline("echo $CONDA_PREFIX") - proc.expect(sys.prefix) - proc.kill(signal.SIGINT) + proc.sendeof() + out = proc.read().decode() + assert "CONDA_SPAWN" in out + assert "CONDA_PREFIX" in out + assert str(simple_env) in out @pytest.mark.skipif(sys.platform != "win32", reason="Powershell only tested on Windows") -def test_powershell(): - shell = PowershellShell() - with shell.spawn_popen(sys.prefix, command=["ls", "env:"], stdout=PIPE, text=True) as proc: - out, _ = proc.communicate() +def test_powershell(simple_env): + shell = PowershellShell(simple_env) + with shell.spawn_popen(command=["ls", "env:"], stdout=PIPE, text=True) as proc: + out, _ = proc.communicate(timeout=5) proc.kill() assert not proc.poll() - assert "CONDA_SPAWN" in out + assert "CONDA_SPAWN" in out assert "CONDA_PREFIX" in out + assert str(simple_env) in out @pytest.mark.skipif(sys.platform != "win32", reason="Cmd.exe only tested on Windows") -def test_cmd(): - shell = CmdExeShell() - with shell.spawn_popen(sys.prefix, command=["@SET"], stdout=PIPE, text=True) as proc: - out, _ = proc.communicate() +def test_cmd(simple_env): + shell = CmdExeShell(simple_env) + with shell.spawn_popen(command=["@SET"], stdout=PIPE, text=True) as proc: + out, _ = proc.communicate(timeout=5) proc.kill() assert not proc.poll() assert "CONDA_SPAWN" in out assert "CONDA_PREFIX" in out + assert str(simple_env) in out + + +def test_hooks(conda_cli, simple_env): + out, err, rc = conda_cli("spawn", "--hook", "-p", simple_env) + print(out) + print(err, file=sys.stderr) + assert not rc + assert not err + assert "CONDA_EXE" in out + assert str(simple_env) in out + + +@pytest.mark.skipif(sys.platform == "win32", reason="Only tested on Unix") +def test_hooks_integration_posix(simple_env, tmp_path): + hook = f"{sys.executable} -m conda spawn --hook --shell posix -p '{simple_env}'" + script = f'eval "$({hook})"\nenv | sort' + script_path = tmp_path / "script-eval.sh" + script_path.write_text(script) + + out = check_output(["bash", script_path], text=True) + print(out) + assert str(simple_env) in out + + +@pytest.mark.skipif(sys.platform != "win32", reason="Powershell only tested on Windows") +def test_hooks_integration_powershell(simple_env, tmp_path): + hook = f"{sys.executable} -m conda spawn --hook --shell powershell -p {simple_env}" + script = f"{hook} | Out-String | Invoke-Expression\r\nls env:" + script_path = tmp_path / "script-eval.ps1" + script_path.write_text(script) + + out = check_output(["powershell", "-NoLogo", "-File", script_path], text=True) + print(out) + assert str(simple_env) in out + + +@pytest.mark.skipif(sys.platform != "win32", reason="Cmd.exe only tested on Windows") +def test_hooks_integration_cmd(simple_env, tmp_path): + hook = f"{sys.executable} -m conda spawn --hook --shell cmd -p {simple_env}" + script = f"FOR /F \"tokens=*\" %%g IN ('{hook}') do @CALL %%g\r\nset" + script_path = tmp_path / "script-eval.bat" + script_path.write_text(script) + + out = check_output(["cmd", "/D", "/C", script_path], text=True) + print(out) + assert str(simple_env) in out