Enable Zsh completions

config/ruff.toml

@@ -53,6 +53,9 @@ ignore = [
 "src/*/debug.py" = [
     "T201",  # Print statement
 ]
+"src/*/completion.py" = [
+    "T201",  # Print statement
+]
 "scripts/*.py" = [
     "INP001",  # File is part of an implicit namespace package
     "T201",  # Print statement

docs/usage.md

@@ -877,8 +877,13 @@ duty task1 task2
 
 ### Shell completions
 
+Duty supports shell completions for Bash and Zsh, these can be automatically installed using:
+```shell
+duty --install-completion
+```
+Completions can also be installed manually:
+
 === "Bash"
-    You can enable auto-completion in Bash with these commands:
     ```bash
     completions_dir="${BASH_COMPLETION_USER_DIR:-${XDG_DATA_HOME:-$HOME/.local/share}/bash-completion}/completions"
     mkdir -p "${completions_dir}"
@@ -895,19 +900,19 @@ duty task1 task2
     ```
 
     If you don't use Oh My Zsh, you can install completions globally under `/usr/local/share/zsh/site-functions`
-    or use a custom directory, for example `~/.duty`.
+    or use a custom directory, for example `~/.zfunc`.
     To do this, make sure that the following get called in your `.zshrc` in this order:
     ```zsh
-    fpath=($HOME/.duty $fpath)
+    fpath=($HOME/.zfunc $fpath)
     autoload -Uz compinit && compinit
     ```
     !!! Warning
         Don't add `autoload -Uz compinit && compinit` when using Oh My Zsh.
 
-    Then generate completion function and restart shell:
+    Then generate completion function and reload shell:
     ```zsh
-    mkdir -p "$HOME/.duty"
-    duty --completion=zsh > "$HOME/.duty/_duty"
+    mkdir -p "$HOME/.zfunc"
+    duty --completion=zsh > "$HOME/.zfunc/_duty"
     exec zsh
     ```
     The completion script file must start with an underscore.
@@ -918,4 +923,4 @@ duty task1 task2
     ```zsh
     autoload -Uz bashcompinit && bashcompinit
     ```
-    Then restart your shell and follow instructions for Bash.
+    Then reload your shell and follow instructions for Bash.

src/duty/cli.py

@@ -18,14 +18,13 @@
 import os
 import sys
 import textwrap
-from pathlib import Path
-from typing import Any
+from typing import Any, Literal
 
 from failprint.cli import ArgParser, add_flags
 
 from duty import debug
-from duty._completion import CompletionParser
 from duty.collection import Collection, Duty
+from duty.completion import CompletionInstaller, CompletionParser
 from duty.exceptions import DutyFailure
 from duty.validation import validate
 
@@ -72,13 +71,21 @@ def get_parser() -> ArgParser:
         metavar="DUTY",
         help="Show this help message and exit. Pass duties names to print their help.",
     )
+    parser.add_argument(
+        "--install-completion",
+        dest="install_completion",
+        nargs="?",
+        const=True,
+        metavar="SHELL",
+        help="Installs completion for the selected shell. If no value is provided, $SHELL is used.",
+    )
     parser.add_argument(
         "--completion",
         dest="completion",
         nargs="?",
         const=True,
         metavar="SHELL",
-        help="Prints completion script for selected shell. If no value is provided, $SHELL is used.",
+        help="Prints completion script for the selected shell. If no value is provided, $SHELL is used.",
     )
     parser.add_argument(
         "--complete",
@@ -261,6 +268,13 @@ def print_help(parser: ArgParser, opts: argparse.Namespace, collection: Collecti
         print(textwrap.indent(collection.format_help(), prefix="  "))
 
 
+def get_shell(arg: str | Literal[True]) -> str:
+    """Get shell from passed arg, or try to guess based on `SHELL` environmental variable."""
+    if arg is True:
+        return os.path.basename(os.environ.get("SHELL", "/bin/bash"))
+    return arg.lower()
+
+
 def main(args: list[str] | None = None) -> int:
     """Run the main program.
 
@@ -279,18 +293,14 @@ def main(args: list[str] | None = None) -> int:
     collection = Collection(opts.duties_file)
     collection.load()
 
-    if opts.completion:
-        if opts.completion is True:
-            shell = os.path.basename(os.environ.get("SHELL", "/bin/bash"))
-        else:
-            shell = opts.completion.lower()
-
-        try:
-            print((Path(__file__).parent / f"completions.{shell}").read_text())
-        except FileNotFoundError as exc:
-            msg = f"Completions for {shell!r} shell are not available, feature requests and PRs welcome!"
-            raise NotImplementedError(msg) from exc
+    if opts.install_completion:
+        shell = get_shell(opts.install_completion)
+        CompletionInstaller.install(shell)
+        return 0
 
+    if opts.completion:
+        shell = get_shell(opts.completion)
+        print(CompletionInstaller.get_completion_script_path(shell).read_text())
         return 0
 
     if opts.complete:

src/duty/collection.py

@@ -12,7 +12,7 @@
 from duty.context import Context
 
 if typing.TYPE_CHECKING:
-    from duty._completion import CompletionCandidateType
+    from duty.completion import CompletionCandidateType
 
 DutyListType = list[Union[str, Callable, "Duty"]]
 default_duties_file = "duties.py"

src/duty/completion.py

@@ -0,0 +1,125 @@
+"""Shell completion utilities."""
+
+import os
+import subprocess
+import sys
+from pathlib import Path
+from typing import Optional
+
+CompletionCandidateType = tuple[str, Optional[str]]
+
+
+class CompletionParser:
+    """Shell completion parser."""
+
+    @classmethod
+    def parse(cls, candidates: list[CompletionCandidateType], shell: str) -> str:
+        """Parses a list of completion candidates for the selected shell's completion command.
+
+        Parameters:
+            candidates: List of completion candidates with optional descriptions.
+            shell: Shell for which to parse the candidates.
+
+        Raises:
+            NotImplementedError: When parser is not implemented for selected shell.
+
+        Returns:
+            String to be passed to shell completion command.
+        """
+        try:
+            return getattr(cls, f"_{shell}")(candidates)
+        except AttributeError as exc:
+            msg = f"Completion parser method for {shell!r} shell is not implemented!"
+            raise NotImplementedError(msg) from exc
+
+    @staticmethod
+    def _zsh(candidates: list[CompletionCandidateType]) -> str:
+        def parse_candidate(item: CompletionCandidateType) -> str:
+            completion, help_text = item
+            # We only have space for one line of description,
+            # so we remove descriptions of sub-command parameters from help_text
+            # by removing everything after the first newline.
+            return f"{completion}: {help_text or '-'}".split("\n", 1)[0]
+
+        return "\n".join(parse_candidate(candidate) for candidate in candidates)
+
+    @staticmethod
+    def _bash(candidates: list[CompletionCandidateType]) -> str:
+        return "\n".join(completion for completion, _ in candidates)
+
+
+class CompletionInstaller:
+    """Shell completion installer."""
+
+    @classmethod
+    def install(cls, shell: str) -> None:
+        """Installs shell completions for selected shell.
+
+        Raises:
+            NotImplementedError: When installer is not implemented for selected shell.
+        """
+        try:
+            return getattr(cls, f"_{shell}")()
+        except AttributeError as exc:
+            msg = f"Completion installer method for {shell!r} shell is not implemented!"
+            raise NotImplementedError(msg) from exc
+
+    @staticmethod
+    def get_completion_script_path(shell: str) -> Path:
+        """Gets the path of a shell completion script for the selected shell."""
+        completions_file_path = Path(__file__).parent / f"completions.{shell}"
+        if not completions_file_path.exists():
+            msg = f"Completions for {shell!r} shell are not available, feature requests and PRs welcome!"
+            raise NotImplementedError(msg)
+        return completions_file_path
+
+    @classmethod
+    def _zsh(cls) -> None:
+        site_functions_dirs = (Path("/usr/local/share/zsh/site-functions"), Path("/usr/share/zsh/site-functions"))
+        try:
+            completions_dir = next(d for d in site_functions_dirs if d.is_dir())
+        except StopIteration as exc:
+            raise OSError("Zsh site-functions directory not found!") from exc
+
+        try:
+            symlink_path = completions_dir / "_duty"
+            symlink_path.symlink_to(cls.get_completion_script_path("zsh"))
+        except PermissionError:
+            # retry as sudo
+            if os.geteuid() == 0:
+                raise
+            subprocess.run(  # noqa: S603
+                ["sudo", sys.executable, sys.argv[0], "--install-completion=zsh"],  # noqa: S607
+                check=True,
+            )
+        except FileExistsError:
+            print("Zsh completions already installed.")
+        else:
+            print(
+                f"Zsh completions successfully symlinked to {symlink_path}. "
+                f"Please reload Zsh for changes to take effect.",
+            )
+
+    @classmethod
+    def _bash(cls) -> None:
+        bash_completion_user_dir = os.environ.get("BASH_COMPLETION_USER_DIR")
+        xdg_data_home = os.environ.get("XDG_DATA_HOME")
+
+        if bash_completion_user_dir:
+            completion_dir = Path(bash_completion_user_dir) / "completions"
+        elif xdg_data_home:
+            completion_dir = Path(xdg_data_home) / "bash-completion/completions"
+        else:
+            completion_dir = Path.home() / ".local/share/bash-completion/completions"
+
+        completion_dir.mkdir(parents=True, exist_ok=True)
+        symlink_path = completion_dir / "duty"
+        try:
+            symlink_path.symlink_to(cls.get_completion_script_path("bash"))
+        except FileExistsError:
+            print("Bash completions already installed.")
+        else:
+            print(
+                f"Bash completions successfully symlinked to {symlink_path!r}. "
+                f"Please reload Bash for changes to take effect.",
+            )