add autocomplete support for zsh and fish shells

.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

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(

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

changelog.d/+zsh_autocomplete.added.md

@@ -0,0 +1 @@
+Add autocomplete support for `zsh` and `fish` shells.

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

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

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.

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)