From f3f554d3d9edfd41a9cabf8c027de7ff072413d1 Mon Sep 17 00:00:00 2001 From: Maciej Urbanski Date: Tue, 26 Mar 2024 19:53:49 +0100 Subject: [PATCH] add autocomplete support for zsh and fish shells --- .github/workflows/ci.yml | 5 + b2/_internal/_cli/autocomplete_install.py | 143 +++++++++++++++--- b2/_internal/_utils/python_compat.py | 5 + changelog.d/+zsh_autocomplete.added.md | 1 + test/unit/_cli/fixtures/dummy_command.py | 24 +++ test/unit/_cli/test_autocomplete_install.py | 34 ++++- test/unit/conftest.py | 13 ++ .../console_tool/test_install_autocomplete.py | 45 ++++++ 8 files changed, 247 insertions(+), 23 deletions(-) create mode 100644 changelog.d/+zsh_autocomplete.added.md create mode 100755 test/unit/_cli/fixtures/dummy_command.py create mode 100644 test/unit/console_tool/test_install_autocomplete.py diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index 0849ff41c..c247af7fb 100644 --- a/.github/workflows/ci.yml +++ b/.github/workflows/ci.yml @@ -95,6 +95,11 @@ jobs: uses: actions/setup-python@v5 with: python-version: ${{ matrix.python-version }} + - name: Install test binary dependencies + if: startsWith(matrix.os, 'ubuntu') + run: | + sudo apt-get -y update + sudo apt-get -y install zsh fish - name: Install dependencies run: python -m pip install --upgrade nox pdm - name: Run unit tests diff --git a/b2/_internal/_cli/autocomplete_install.py b/b2/_internal/_cli/autocomplete_install.py index 562ec2494..6390a52aa 100644 --- a/b2/_internal/_cli/autocomplete_install.py +++ b/b2/_internal/_cli/autocomplete_install.py @@ -7,17 +7,24 @@ # License https://www.backblaze.com/using_b2_code.html # ###################################################################### + +import abc import logging +import platform import re import shutil import subprocess +import textwrap from datetime import datetime from pathlib import Path +from shlex import quote from typing import List import argcomplete from class_registry import ClassRegistry, RegistryKeyError +from b2._internal._utils.python_compat import shlex_join + logger = logging.getLogger(__name__) SHELL_REGISTRY = ClassRegistry() @@ -33,7 +40,7 @@ def autocomplete_install(prog: str, shell: str = 'bash') -> None: logger.info("Autocomplete for %s has been enabled.", prog) -class ShellAutocompleteInstaller: +class ShellAutocompleteInstaller(abc.ABC): shell_exec: str def __init__(self, prog: str): @@ -43,6 +50,10 @@ def install(self) -> None: """Install autocomplete for the given program.""" script_path = self.create_script() if not self.is_enabled(): + logger.info( + "%s completion doesn't seem to be autoloaded from %s.", self.shell_exec, + script_path.parent + ) try: self.force_enable(script_path) except NotImplementedError as e: @@ -60,10 +71,11 @@ def create_script(self) -> Path: script_path = self.get_script_path() logger.info("Creating autocompletion script under %s", script_path) - script_path.parent.mkdir(exist_ok=True) + script_path.parent.mkdir(exist_ok=True, parents=True) script_path.write_text(shellcode) return script_path + @abc.abstractmethod def force_enable(self, completion_script: Path) -> None: """ Enable autocomplete for the given program. @@ -76,6 +88,7 @@ def get_shellcode(self) -> str: """Get autocomplete shellcode for the given program.""" return argcomplete.shellcode([self.prog], shell=self.shell_exec) + @abc.abstractmethod def get_script_path(self) -> Path: """Get autocomplete script path for the given program.""" raise NotImplementedError @@ -84,43 +97,129 @@ def program_in_path(self) -> bool: """Check if the given program is in PATH.""" return _silent_success_run([self.shell_exec, '-c', self.prog]) + @abc.abstractmethod def is_enabled(self) -> bool: """Check if autocompletion is enabled.""" - return _silent_success_run([self.shell_exec, '-i', '-c', f'complete -p {self.prog}']) + raise NotImplementedError -@SHELL_REGISTRY.register('bash') -class BashAutocompleteInstaller(ShellAutocompleteInstaller): - shell_exec = 'bash' +class BashLikeAutocompleteInstaller(ShellAutocompleteInstaller): + shell_exec: str + rc_file_path: str def force_enable(self, completion_script: Path) -> None: - """Enable autocomplete for the given program.""" - logger.info( - "Bash completion doesn't seem to be autoloaded from %s. Most likely `bash-completion` is not installed.", - completion_script.parent - ) - bashrc_path = Path("~/.bashrc").expanduser() - if bashrc_path.exists(): - bck_path = bashrc_path.with_suffix(f".{datetime.now():%Y-%m-%dT%H-%M-%S}.bak") - logger.warning("Backing up %s to %s", bashrc_path, bck_path) + """Enable autocomplete for the given program, common logic.""" + rc_path = Path(self.rc_file_path).expanduser() + if rc_path.exists(): + bck_path = rc_path.with_suffix(f".{datetime.now():%Y-%m-%dT%H-%M-%S}.bak") + logger.warning("Backing up %s to %s", rc_path, bck_path) try: - shutil.copyfile(bashrc_path, bck_path) + shutil.copyfile(rc_path, bck_path) except OSError as e: raise AutocompleteInstallError( - f"Failed to backup {bashrc_path} under {bck_path}" + f"Failed to backup {rc_path} under {bck_path}" ) from e - logger.warning("Explicitly adding %s to %s", completion_script, bashrc_path) + logger.warning("Explicitly adding %s to %s", completion_script, rc_path) add_or_update_shell_section( - bashrc_path, f"{self.prog} autocomplete", self.prog, f"source {completion_script}" + rc_path, f"{self.prog} autocomplete", self.prog, self.get_rc_section(completion_script) ) + def get_rc_section(self, completion_script: Path) -> str: + return f"source {quote(str(completion_script))}" + def get_script_path(self) -> Path: - """Get autocomplete script path for the given program.""" - return Path("~/.bash_completion.d/").expanduser() / self.prog + """Get autocomplete script path for the given program, common logic.""" + script_dir = Path(f"~/.{self.shell_exec}_completion.d/").expanduser() + return script_dir / self.prog + + def is_enabled(self) -> bool: + """Check if autocompletion is enabled.""" + return _silent_success_run([self.shell_exec, '-i', '-c', f'complete -p {quote(self.prog)}']) + + +@SHELL_REGISTRY.register('bash') +class BashAutocompleteInstaller(BashLikeAutocompleteInstaller): + shell_exec = 'bash' + rc_file_path = "~/.bashrc" + + +@SHELL_REGISTRY.register('zsh') +class ZshAutocompleteInstaller(BashLikeAutocompleteInstaller): + shell_exec = 'zsh' + rc_file_path = "~/.zshrc" + + def get_rc_section(self, completion_script: Path) -> str: + return textwrap.dedent( + f"""\ + if [[ -z "$_comps" ]] && [[ -t 0 ]]; then autoload -Uz compinit && compinit; fi + source {quote(str(completion_script))} + """ + ) + + def get_script_path(self) -> Path: + """Custom get_script_path for Zsh, if the structure differs from the base implementation.""" + return Path("~/.zsh/completion/").expanduser() / f"_{self.prog}" + + def is_enabled(self) -> bool: + # fallback to tty emulation since `zsh` fails without tty (e.g. under GitHub Actions) + cmd = [self.shell_exec, '-i', '-c', f'[[ -v _comps[{quote(self.prog)}] ]]'] + return _silent_success_run(cmd) or _silent_success_run(_tty_cmd_wrap(cmd)) + + +@SHELL_REGISTRY.register('fish') +class FishAutocompleteInstaller(ShellAutocompleteInstaller): + shell_exec = 'fish' + rc_file_path = "~/.config/fish/config.fish" + + def force_enable(self, completion_script: Path) -> None: + raise NotImplementedError("Fish shell doesn't support manual completion enabling.") + + def get_script_path(self) -> Path: + """Get autocomplete script path for the given program, common logic.""" + return Path("~/.config/fish/completions/").expanduser() / f"{self.prog}.fish" + + def is_enabled(self) -> bool: + """ + Check if autocompletion is enabled. + + Fish seems to lazy-load completions, hence first we trigger completion. + That alone cannot be used, since fish tends to always propose completions (e.g. suggesting similarly + named filenames). + """ + return _silent_success_run( + [ + self.shell_exec, '-i', '-c', + f'string length -q -- (complete -C{quote(f"{self.prog} ")} >/dev/null && complete -c {quote(self.prog)})' + ] + ) def _silent_success_run(cmd: List[str]) -> bool: - return subprocess.run(cmd, stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL).returncode == 0 + p = subprocess.Popen( + cmd, stdout=subprocess.PIPE, stderr=subprocess.PIPE, stdin=subprocess.DEVNULL + ) + stdout, stderr = p.communicate() + if p.returncode != 0: + logger.debug( + "Command %s exited with code %d, stdout: %s, stderr: %s", cmd, p.returncode, stdout, + stderr + ) + return p.returncode == 0 + + +def _tty_cmd_wrap(cmd: List[str]) -> List[str]: + """ + Add tty emulation to the command. + + Inspired by tty emulation for different platforms found here: + https://github.com/Yuri6037/Action-FakeTTY/tree/master/script + """ + platform_name = platform.platform().lower() + if platform_name == 'darwin': + return ['unbuffer', shlex_join(cmd)] + elif platform_name == 'windows': + return ["Invoke-Expression", shlex_join(cmd)] + return ['script', '-q', '-e', '-c', shlex_join(cmd), '/dev/null'] def add_or_update_shell_section( diff --git a/b2/_internal/_utils/python_compat.py b/b2/_internal/_utils/python_compat.py index 54b82143d..59553e62d 100644 --- a/b2/_internal/_utils/python_compat.py +++ b/b2/_internal/_utils/python_compat.py @@ -11,6 +11,7 @@ Utilities for compatibility with older Python versions. """ import functools +import shlex import sys if sys.version_info < (3, 9): @@ -45,5 +46,9 @@ def method_wrapper(arg, *args, **kwargs): method_wrapper.register = self.register return method_wrapper + + def shlex_join(split_command): + return ' '.join(shlex.quote(arg) for arg in split_command) else: singledispatchmethod = functools.singledispatchmethod + shlex_join = shlex.join diff --git a/changelog.d/+zsh_autocomplete.added.md b/changelog.d/+zsh_autocomplete.added.md new file mode 100644 index 000000000..6cfd3cfba --- /dev/null +++ b/changelog.d/+zsh_autocomplete.added.md @@ -0,0 +1 @@ +Add autocomplete support for `zsh` and `fish` shells. diff --git a/test/unit/_cli/fixtures/dummy_command.py b/test/unit/_cli/fixtures/dummy_command.py new file mode 100755 index 000000000..82ec4b76f --- /dev/null +++ b/test/unit/_cli/fixtures/dummy_command.py @@ -0,0 +1,24 @@ +#!/usr/bin/env python +###################################################################### +# +# File: test/unit/_cli/fixtures/dummy_command.py +# +# Copyright 2024 Backblaze Inc. All Rights Reserved. +# +# License https://www.backblaze.com/using_b2_code.html +# +###################################################################### +import argparse + + +def main(): + parser = argparse.ArgumentParser(description="Dummy command") + parser.add_argument("--foo", help="foo help") + parser.add_argument("--bar", help="bar help") + args = parser.parse_args() + print(args.foo) + print(args.bar) + + +if __name__ == "__main__": + main() diff --git a/test/unit/_cli/test_autocomplete_install.py b/test/unit/_cli/test_autocomplete_install.py index f6dfdd64b..50b76cc50 100644 --- a/test/unit/_cli/test_autocomplete_install.py +++ b/test/unit/_cli/test_autocomplete_install.py @@ -7,9 +7,16 @@ # License https://www.backblaze.com/using_b2_code.html # ###################################################################### +import pathlib +import shutil +from test.helpers import skip_on_windows + import pytest -from b2._internal._cli.autocomplete_install import add_or_update_shell_section +from b2._internal._cli.autocomplete_install import ( + SHELL_REGISTRY, + add_or_update_shell_section, +) section = "test_section" managed_by = "pytest" @@ -73,3 +80,28 @@ def test_add_or_update_shell_section_no_file(test_file): {content} # <<< {section} <<< """ + + +@pytest.fixture +def dummy_command(homedir, monkeypatch, env): + name = "dummy_command" + bin_path = homedir / "bin" / name + bin_path.parent.mkdir(parents=True, exist_ok=True) + bin_path.symlink_to(pathlib.Path(__file__).parent / "fixtures" / f"{name}.py") + monkeypatch.setenv("PATH", f"{homedir}/bin:{env['PATH']}") + yield name + + +@pytest.mark.parametrize("shell", ["bash", "zsh", "fish"]) +@skip_on_windows +def test_autocomplete_installer(homedir, env, shell, caplog, dummy_command): + caplog.set_level(10) + shell_installer = SHELL_REGISTRY.get(shell, prog=dummy_command) + + shell_bin = shutil.which(shell) + if shell_bin is None: + pytest.skip(f"{shell} is not installed") + + assert shell_installer.is_enabled() is False + shell_installer.install() + assert shell_installer.is_enabled() is True diff --git a/test/unit/conftest.py b/test/unit/conftest.py index 53cb5004e..8719e9268 100644 --- a/test/unit/conftest.py +++ b/test/unit/conftest.py @@ -43,6 +43,19 @@ def cli_version(request) -> str: return request.config.getoption('--cli') +@pytest.fixture +def homedir(tmp_path_factory): + yield tmp_path_factory.mktemp("test_homedir") + + +@pytest.fixture +def env(homedir, monkeypatch): + """Get ENV for running b2 command from shell level.""" + monkeypatch.setenv("HOME", str(homedir)) + monkeypatch.setenv("SHELL", "/bin/bash") # fix for running under github actions + yield os.environ + + @pytest.fixture(scope='session') def console_tool_class(cli_version): # Ensures import of the correct library to handle all the tests. diff --git a/test/unit/console_tool/test_install_autocomplete.py b/test/unit/console_tool/test_install_autocomplete.py new file mode 100644 index 000000000..3f3caa3df --- /dev/null +++ b/test/unit/console_tool/test_install_autocomplete.py @@ -0,0 +1,45 @@ +###################################################################### +# +# File: test/unit/console_tool/test_install_autocomplete.py +# +# Copyright 2024 Backblaze Inc. All Rights Reserved. +# +# License https://www.backblaze.com/using_b2_code.html +# +###################################################################### + +import contextlib +import shutil +from test.helpers import skip_on_windows + +import pexpect +import pytest + + +@contextlib.contextmanager +def pexpect_shell(shell_bin, env): + p = pexpect.spawn(f"{shell_bin} -i", env=env, maxread=1000) + p.setwinsize(100, 100) # required to see all suggestions in tests + yield p + p.close() + + +@pytest.mark.parametrize("shell", ["bash", "zsh", "fish"]) +@skip_on_windows +def test_install_autocomplete(b2_cli, env, shell, monkeypatch): + shell_bin = shutil.which(shell) + if shell_bin is None: + pytest.skip(f"{shell} is not installed") + + monkeypatch.setenv("SHELL", shell_bin) + b2_cli.run( + ["install-autocomplete"], + expected_part_of_stdout=f"Autocomplete successfully installed for {shell}", + ) + + if shell == "fish": # no idea how to test fish autocompletion (does not seem to work with dummy terminal) + return + + with pexpect_shell(shell_bin, env=env) as pshell: + pshell.send("b2 \t\t") + pshell.expect_exact(["authorize-account", "download-file", "get-bucket"], timeout=30)