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,12 +7,14 @@
 # License https://www.backblaze.com/using_b2_code.html
 #
 ######################################################################
+import abc
 import logging
 import re
 import shutil
 import subprocess
 from datetime import datetime
 from pathlib import Path
+from shlex import quote
 from typing import List
 
 import argcomplete
@@ -33,7 +35,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 +45,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 +66,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 +83,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,39 +92,85 @@ 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, f"source {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_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:
+        return _silent_success_run(
+            [self.shell_exec, '-i', '-c', f'[[ -v _comps[{quote(self.prog)}] ]]']
+        )
+
+
+@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.
+
+        There seems to be no reliable way to check if a completion is enabled in fish.
+        It seems like completion is loaded (asynchronously?) some time after launching the shell.
+        `complete -c b2` works, but not immediately and not when run using `fish -i -c`.
+        """
+        return self.get_script_path().exists()
 
 
 def _silent_success_run(cmd: List[str]) -> bool:

changelog.d/+zsh_autocomplete.added.md

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

test/unit/_cli/test_autocomplete_install.py

@@ -7,9 +7,15 @@
 # License https://www.backblaze.com/using_b2_code.html
 #
 ######################################################################
+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 +79,25 @@ def test_add_or_update_shell_section_no_file(test_file):
 {content}
 # <<< {section} <<<
 """
+
+
+@pytest.mark.parametrize(
+    "shell",
+    [
+        "bash",
+        "zsh",
+        "fish",
+    ],
+)
+@skip_on_windows
+def test_autocomplete_installer(homedir, env, shell):
+    shell_bin = shutil.which(shell)
+    if shell_bin is None:
+        pytest.skip(f"{shell} is not installed")
+
+    cmd_name = "dummy_command"
+    shell_installer = SHELL_REGISTRY.get(shell, prog=cmd_name)
+
+    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,52 @@
+######################################################################
+#
+# 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)