From 5c52138f38d4b4662a8f8fb0110ddb5814ab152f Mon Sep 17 00:00:00 2001 From: Lam Chau Date: Wed, 25 Sep 2024 15:58:19 -0700 Subject: [PATCH] feat: add shell-completions subcommand (#76) --- pyproject.toml | 25 ++++---- src/goose/cli/main.py | 59 +++++++++++++++++-- src/goose/utils/autocomplete.py | 100 ++++++++++++++++++++++++++++++++ tests/test_autocomplete.py | 34 +++++++++++ tests/test_cli_main.py | 21 +++++++ 5 files changed, 221 insertions(+), 18 deletions(-) create mode 100644 src/goose/utils/autocomplete.py create mode 100644 tests/test_autocomplete.py create mode 100644 tests/test_cli_main.py diff --git a/pyproject.toml b/pyproject.toml index ba4388e4..8913104c 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -48,20 +48,21 @@ build-backend = "hatchling.build" [tool.uv] dev-dependencies = [ - "pytest>=8.3.2", "codecov>=2.1.13", - "mkdocstrings>=0.26.1", - "mkdocs-literate-nav>=0.6.1", - "mkdocs-gen-files>=0.5.0", - "mkdocs-section-index>=0.3.9", - "mkdocs-material>=9.5.34", - "mkdocstrings-python>=1.11.1", - "mkdocs-git-revision-date-localized-plugin", - "mkdocs-glightbox>=0.4.0", - "mkdocs-redirects>=1.2.1", - "mkdocs-include-markdown-plugin>=6.2.2", "mkdocs-callouts>=1.14.0", + "mkdocs-gen-files>=0.5.0", "mkdocs-git-authors-plugin>=0.9.0", - "mkdocs-git-revision-date-localized-plugin>=1.2.9", "mkdocs-git-committers-plugin>=0.2.3", + "mkdocs-git-revision-date-localized-plugin", + "mkdocs-git-revision-date-localized-plugin>=1.2.9", + "mkdocs-glightbox>=0.4.0", + "mkdocs-include-markdown-plugin>=6.2.2", + "mkdocs-literate-nav>=0.6.1", + "mkdocs-material>=9.5.34", + "mkdocs-redirects>=1.2.1", + "mkdocs-section-index>=0.3.9", + "mkdocstrings-python>=1.11.1", + "mkdocstrings>=0.26.1", + "pytest-mock>=3.14.0", + "pytest>=8.3.2" ] diff --git a/src/goose/cli/main.py b/src/goose/cli/main.py index 5d3a6b31..a0495a30 100644 --- a/src/goose/cli/main.py +++ b/src/goose/cli/main.py @@ -1,3 +1,4 @@ +import os from datetime import datetime from pathlib import Path from typing import Optional @@ -10,6 +11,7 @@ from goose.cli.config import SESSIONS_PATH from goose.cli.session import Session from goose.toolkit.utils import render_template, parse_plan from goose.utils import load_plugins +from goose.utils.autocomplete import SUPPORTED_SHELLS, setup_autocomplete from goose.utils.session_file import list_sorted_session_files @@ -43,6 +45,38 @@ def get_version() -> None: print(f" [red]Could not retrieve version for {module}: {e}[/red]") +def get_current_shell() -> str: + return os.getenv("SHELL", "").split("/")[-1] + + +@goose_cli.command(name="shell-completions", help="Manage shell completions for goose") +@click.option("--install", is_flag=True, help="Install shell completions") +@click.option("--generate", is_flag=True, help="Generate shell completions") +@click.argument( + "shell", + type=click.Choice(SUPPORTED_SHELLS), + default=get_current_shell(), +) +@click.pass_context +def shell_completions(ctx: click.Context, install: bool, generate: bool, shell: str) -> None: + """Generate or install shell completions for goose + + Args: + shell (str): shell to install completions for + install (bool): installs completions if true, otherwise generates + completions + """ + if not any([install, generate]): + print("[red]One of --install or --generate must be specified[/red]\n") + raise click.UsageError(ctx.get_help()) + + if sum([install, generate]) > 1: + print("[red]Only one of --install or --generate can be specified[/red]\n") + raise click.UsageError(ctx.get_help()) + + setup_autocomplete(shell, install=install) + + @goose_cli.group() def session() -> None: """Start or manage sessions""" @@ -100,19 +134,36 @@ def session_planned(plan: str, args: Optional[dict[str, str]]) -> None: session.run() +def autocomplete_session_files(ctx: click.Context, args: str, incomplete: str) -> None: + return [ + f"{session_name}" + for session_name in sorted(get_session_files().keys(), reverse=True, key=lambda x: x.lower()) + if session_name.startswith(incomplete) + ] + + +def get_session_files() -> dict[str, Path]: + return list_sorted_session_files(SESSIONS_PATH) + + @session.command(name="resume") -@click.argument("name", required=False) +@click.argument("name", required=False, shell_complete=autocomplete_session_files) @click.option("--profile") def session_resume(name: Optional[str], profile: str) -> None: """Resume an existing goose session""" + session_files = get_session_files() if name is None: - session_files = get_session_files() if session_files: name = list(session_files.keys())[0] print(f"Resuming most recent session: {name} from {session_files[name]}") else: print("No sessions found.") return + else: + if name in session_files: + print(f"Resuming session: {name}") + else: + print(f"Creating new session: {name}") session = Session(name=name, profile=profile) session.run() @@ -134,10 +185,6 @@ def session_clear(keep: int) -> None: session_file.unlink() -def get_session_files() -> dict[str, Path]: - return list_sorted_session_files(SESSIONS_PATH) - - @click.group( invoke_without_command=True, name="goose", diff --git a/src/goose/utils/autocomplete.py b/src/goose/utils/autocomplete.py new file mode 100644 index 00000000..6feb0807 --- /dev/null +++ b/src/goose/utils/autocomplete.py @@ -0,0 +1,100 @@ +import sys +from pathlib import Path + +from rich import print + +SUPPORTED_SHELLS = ["bash", "zsh", "fish"] + + +def is_autocomplete_installed(file: Path) -> bool: + if not file.exists(): + print(f"[yellow]{file} does not exist, creating file") + with open(file, "w") as f: + f.write("") + + # https://click.palletsprojects.com/en/8.1.x/shell-completion/#enabling-completion + if "_GOOSE_COMPLETE" in open(file).read(): + print(f"auto-completion already installed in {file}") + return True + return False + + +def setup_bash(install: bool) -> None: + bashrc = Path("~/.bashrc").expanduser() + if install: + if is_autocomplete_installed(bashrc): + return + f = open(bashrc, "a") + else: + f = sys.stdout + print(f"# add the following to your bash config, typically {bashrc}") + + with f: + f.write('eval "$(_GOOSE_COMPLETE=bash_source goose)"\n') + + if install: + print(f"installed auto-completion to {bashrc}") + print(f"run `source {bashrc}` to enable auto-completion") + + +def setup_fish(install: bool) -> None: + completion_dir = Path("~/.config/fish/completions").expanduser() + if not completion_dir.exists(): + completion_dir.mkdir(parents=True, exist_ok=True) + + completion_file = completion_dir / "goose.fish" + if install: + if is_autocomplete_installed(completion_file): + return + f = open(completion_file, "a") + else: + f = sys.stdout + print(f"# add the following to your fish config, typically {completion_file}") + + with f: + f.write("_GOOSE_COMPLETE=fish_source goose | source\n") + + if install: + print(f"installed auto-completion to {completion_file}") + + +def setup_zsh(install: bool) -> None: + zshrc = Path("~/.zshrc").expanduser() + if install: + if is_autocomplete_installed(zshrc): + return + f = open(zshrc, "a") + else: + f = sys.stdout + print(f"# add the following to your zsh config, typically {zshrc}") + + with f: + f.write("autoload -U +X compinit && compinit\n") + f.write("autoload -U +X bashcompinit && bashcompinit\n") + f.write('eval "$(_GOOSE_COMPLETE=zsh_source goose)"\n') + + if install: + print(f"installed auto-completion to {zshrc}") + print(f"run `source {zshrc}` to enable auto-completion") + + +def setup_autocomplete(shell: str, install: bool) -> None: + """Installs shell completions for goose + + Args: + shell (str): shell to install completions for + install (bool): whether to install or generate completions + """ + + match shell: + case "bash": + setup_bash(install=install) + + case "zsh": + setup_zsh(install=install) + + case "fish": + setup_fish(install=install) + + case _: + print(f"Shell {shell} not supported") diff --git a/tests/test_autocomplete.py b/tests/test_autocomplete.py new file mode 100644 index 00000000..789b5ec2 --- /dev/null +++ b/tests/test_autocomplete.py @@ -0,0 +1,34 @@ +import sys +import unittest.mock as mock + +from goose.utils.autocomplete import SUPPORTED_SHELLS, is_autocomplete_installed, setup_autocomplete + + +def test_supported_shells(): + assert SUPPORTED_SHELLS == ["bash", "zsh", "fish"] + + +def test_install_autocomplete(tmp_path): + file = tmp_path / "test_bash_autocomplete" + assert is_autocomplete_installed(file) is False + + file.write_text("_GOOSE_COMPLETE") + assert is_autocomplete_installed(file) is True + + +@mock.patch("sys.stdout") +def test_setup_bash(mocker: mock.MagicMock): + setup_autocomplete("bash", install=False) + sys.stdout.write.assert_called_with('eval "$(_GOOSE_COMPLETE=bash_source goose)"\n') + + +@mock.patch("sys.stdout") +def test_setup_zsh(mocker: mock.MagicMock): + setup_autocomplete("zsh", install=False) + sys.stdout.write.assert_called_with('eval "$(_GOOSE_COMPLETE=zsh_source goose)"\n') + + +@mock.patch("sys.stdout") +def test_setup_fish(mocker: mock.MagicMock): + setup_autocomplete("fish", install=False) + sys.stdout.write.assert_called_with("_GOOSE_COMPLETE=fish_source goose | source\n") diff --git a/tests/test_cli_main.py b/tests/test_cli_main.py new file mode 100644 index 00000000..9160be1b --- /dev/null +++ b/tests/test_cli_main.py @@ -0,0 +1,21 @@ +from click.testing import CliRunner +from goose.cli.main import get_current_shell, shell_completions + + +def test_get_current_shell(mocker): + mocker.patch("os.getenv", return_value="/bin/bash") + assert get_current_shell() == "bash" + + +def test_shell_completions_install_invalid_combination(): + runner = CliRunner() + result = runner.invoke(shell_completions, ["--install", "--generate", "bash"]) + assert result.exit_code != 0 + assert "Only one of --install or --generate can be specified" in result.output + + +def test_shell_completions_install_no_option(): + runner = CliRunner() + result = runner.invoke(shell_completions, ["bash"]) + assert result.exit_code != 0 + assert "One of --install or --generate must be specified" in result.output