Skip to content

Commit

Permalink
feat: add shell-completions subcommand (#76)
Browse files Browse the repository at this point in the history
  • Loading branch information
lamchau authored Sep 25, 2024
1 parent 2b6db66 commit 5c52138
Show file tree
Hide file tree
Showing 5 changed files with 220 additions and 17 deletions.
23 changes: 12 additions & 11 deletions pyproject.toml
Original file line number Diff line number Diff line change
Expand Up @@ -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-callouts>=1.14.0",
"mkdocs-gen-files>=0.5.0",
"mkdocs-section-index>=0.3.9",
"mkdocs-material>=9.5.34",
"mkdocstrings-python>=1.11.1",
"mkdocs-git-authors-plugin>=0.9.0",
"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-redirects>=1.2.1",
"mkdocs-include-markdown-plugin>=6.2.2",
"mkdocs-callouts>=1.14.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-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"
]
59 changes: 53 additions & 6 deletions src/goose/cli/main.py
Original file line number Diff line number Diff line change
@@ -1,3 +1,4 @@
import os
from datetime import datetime
from pathlib import Path
from typing import Optional
Expand All @@ -10,6 +11,7 @@
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


Expand Down Expand Up @@ -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"""
Expand Down Expand Up @@ -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()

Expand All @@ -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",
Expand Down
100 changes: 100 additions & 0 deletions src/goose/utils/autocomplete.py
Original file line number Diff line number Diff line change
@@ -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")
34 changes: 34 additions & 0 deletions tests/test_autocomplete.py
Original file line number Diff line number Diff line change
@@ -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")
21 changes: 21 additions & 0 deletions tests/test_cli_main.py
Original file line number Diff line number Diff line change
@@ -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

0 comments on commit 5c52138

Please sign in to comment.