From f5caea992990ee623183f7ec18cd2c797a624f99 Mon Sep 17 00:00:00 2001 From: Oxygen <1391083091@qq.com> Date: Sat, 6 Jun 2026 00:53:46 +0800 Subject: [PATCH] feat: add shell auto-completion for bash, zsh, fish, powershell Adds 'openai completion' subcommand that generates shell completion scripts for supported shells. Closes #843 Co-Authored-By: Claude Opus 4.8 --- pyproject.toml | 3 + src/openai/__main__.py | 3 + src/openai/cli/__init__.py | 1 + src/openai/cli/_cli.py | 300 +++++++++++++++++++++++++++++++++++++ tests/test_cli.py | 85 +++++++++++ 5 files changed, 392 insertions(+) create mode 100644 src/openai/__main__.py create mode 100644 src/openai/cli/__init__.py create mode 100644 src/openai/cli/_cli.py create mode 100644 tests/test_cli.py diff --git a/pyproject.toml b/pyproject.toml index d3a6a58a76..9ddbd7cda0 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -48,6 +48,9 @@ realtime = ["websockets >= 13, < 16"] datalib = ["numpy >= 1", "pandas >= 1.2.3", "pandas-stubs >= 1.1.0.11"] voice_helpers = ["sounddevice>=0.5.1", "numpy>=2.0.2"] +[project.scripts] +openai = "openai.cli:main" + [tool.rye] managed = true # version pins are in requirements-dev.lock diff --git a/src/openai/__main__.py b/src/openai/__main__.py new file mode 100644 index 0000000000..4e28416e10 --- /dev/null +++ b/src/openai/__main__.py @@ -0,0 +1,3 @@ +from .cli import main + +main() diff --git a/src/openai/cli/__init__.py b/src/openai/cli/__init__.py new file mode 100644 index 0000000000..d453d5e179 --- /dev/null +++ b/src/openai/cli/__init__.py @@ -0,0 +1 @@ +from ._cli import main as main diff --git a/src/openai/cli/_cli.py b/src/openai/cli/_cli.py new file mode 100644 index 0000000000..74c9c09e42 --- /dev/null +++ b/src/openai/cli/_cli.py @@ -0,0 +1,300 @@ +from __future__ import annotations + +import sys +import argparse +from typing import Optional + +from .. import __version__ + +_SHELL_TYPES = ("bash", "zsh", "fish", "powershell") + +_BASH_COMPLETION_SCRIPT = """# openai completion for bash +# Save this to a file and source it, or add to ~/.bashrc: +# eval "$(openai completion -s bash)" + +_openai_completion() { + local cur prev words cword + _init_completion || return + + # Collect all words up to current position for context + COMPREPLY=() + cur="${COMP_WORDS[COMP_CWORD]}" + prev="${COMP_WORDS[COMP_CWORD-1]}" + + # Complete using the CLI itself + local IFS=$'\\n' + local completions + completions="$(OPENAI_COMPLETE=bash_source openai "${COMP_WORDS[@]:1}" 2>/dev/null)" + if [[ $? -eq 0 ]]; then + COMPREPLY=($(compgen -W "$completions" -- "$cur")) + fi + return 0 +} + +complete -F _openai_completion openai +""" + +_ZSH_COMPLETION_SCRIPT = """#compdef openai +# openai completion for zsh +# Save this to a file in your fpath, or add to ~/.zshrc: +# eval "$(openai completion -s zsh)" + +_openai() { + local -a completions + local curcontext="$curcontext" state line + typeset -A opt_args + + local IFS=$'\\n' + completions=("${(@f)$(OPENAI_COMPLETE=zsh_source openai "${words[@]:1}" 2>/dev/null)}") + + if [[ -n $completions ]]; then + _describe 'values' completions + fi +} + +_openai +""" + +_FISH_COMPLETION_SCRIPT = """# openai completion for fish +# Save this to ~/.config/fish/completions/openai.fish +# or generate with: openai completion -s fish > ~/.config/fish/completions/openai.fish + +function _openai_completion + set -l args (commandline -opc) + set -l current (commandline -ct) + set -l completions (OPENAI_COMPLETE=fish_source openai $args 2>/dev/null) + if test $status -eq 0 + for comp in $completions + echo $comp + end + end +end + +complete -f -c openai -a "(_openai_completion)" +""" + +_POWERSHELL_COMPLETION_SCRIPT = """# openai completion for PowerShell +# Add to your PowerShell profile: +# openai completion -s powershell | Out-String | Invoke-Expression + +Register-ArgumentCompleter -Native -CommandName openai -ScriptBlock { + param($wordToComplete, $commandAst, $cursorPosition) + + $commandElements = $commandAst.CommandElements + $arguments = @() + for ($i = 1; $i -lt $commandElements.Count; $i++) { + $arguments += $commandElements[$i].Value + } + + $env:OPENAI_COMPLETE = "powershell_source" + $result = & openai $arguments 2>$null + + if ($LASTEXITCODE -eq 0) { + $result | ForEach-Object { + [System.Management.Automation.CompletionResult]::new($_, $_, 'ParameterValue', $_) + } + } +} +""" + + +def _get_completions(subcommands: list[str], flags: list[str], current_word: str) -> list[str]: + """Return possible completions based on the current word.""" + all_options = subcommands + flags + if not current_word: + return all_options + return [opt for opt in all_options if opt.startswith(current_word)] + + +def _handle_bash_source(args: list[str]) -> None: + """Handle bash completion source requests.""" + # Parse the command line to understand context + subcommands = ["completion", "--help", "-h", "--version", "-V"] + flags = ["--help", "-h", "--version", "-V"] + + if len(args) <= 1: + # Completing the first argument + for opt in subcommands: + if not args or (len(args) == 1 and opt.startswith(args[0])): + print(opt) + else: + current = args[-1] if args[-1] != "" else "" + prev = args[-2] if len(args) >= 2 else "" + + if prev == "completion": + # Complete completion subcommand flags + completion_opts = ["--help", "-h", "-s", "--shell"] + for opt in completion_opts: + if opt.startswith(current): + print(opt) + elif prev in ("-s", "--shell"): + # Complete shell types + for shell in _SHELL_TYPES: + if shell.startswith(current): + print(shell) + + +def _handle_zsh_source(args: list[str]) -> None: + """Handle zsh completion source requests.""" + subcommands = ["completion", "--help", "-h", "--version", "-V"] + flags = ["--help", "-h", "--version", "-V"] + + if len(args) <= 1: + for opt in subcommands: + if not args or (len(args) == 1 and opt.startswith(args[0])): + print(opt) + else: + current = args[-1] if args[-1] != "" else "" + prev = args[-2] if len(args) >= 2 else "" + + if prev == "completion": + completion_opts = ["--help", "-h", "-s", "--shell"] + for opt in completion_opts: + if opt.startswith(current): + print(opt) + elif prev in ("-s", "--shell"): + for shell in _SHELL_TYPES: + if shell.startswith(current): + print(shell) + + +def _handle_fish_source(args: list[str]) -> None: + """Handle fish completion source requests.""" + subcommands = ["completion", "--help", "-h", "--version", "-V"] + + if len(args) <= 1: + for opt in subcommands: + if not args or (len(args) == 1 and opt.startswith(args[0])): + print(opt) + else: + current = args[-1] if args[-1] != "" else "" + prev = args[-2] if len(args) >= 2 else "" + + if prev == "completion": + completion_opts = ["--help", "-h", "-s", "--shell"] + for opt in completion_opts: + if opt.startswith(current): + print(opt) + elif prev in ("-s", "--shell"): + for shell in _SHELL_TYPES: + if shell.startswith(current): + print(shell) + + +def _handle_powershell_source(args: list[str]) -> None: + """Handle PowerShell completion source requests.""" + subcommands = ["completion", "--help", "-h", "--version", "-V"] + + if len(args) <= 1: + for opt in subcommands: + if not args or (len(args) == 1 and opt.startswith(args[0])): + print(opt) + else: + current = args[-1] if args[-1] != "" else "" + prev = args[-2] if len(args) >= 2 else "" + + if prev == "completion": + completion_opts = ["--help", "-h", "-s", "--shell"] + for opt in completion_opts: + if opt.startswith(current): + print(opt) + elif prev in ("-s", "--shell"): + for shell in _SHELL_TYPES: + if shell.startswith(current): + print(shell) + + +def _build_parser() -> argparse.ArgumentParser: + parser = argparse.ArgumentParser( + description="OpenAI CLI", + prog="openai", + ) + parser.add_argument( + "-V", + "--version", + action="version", + version="%(prog)s " + __version__, + ) + + subparsers = parser.add_subparsers(dest="command") + + # completion subcommand + completion_parser = subparsers.add_parser( + "completion", + help="Generate shell completion scripts", + description="Generate shell completion scripts for the openai CLI.", + ) + completion_parser.add_argument( + "-s", + "--shell", + type=str, + choices=_SHELL_TYPES, + required=True, + help="Shell type to generate completion for", + ) + completion_parser.set_defaults(func=_run_completion) + + return parser + + +def _run_completion(args: argparse.Namespace) -> None: + """Generate and print the completion script for the specified shell.""" + shell = args.shell + + if shell == "bash": + print(_BASH_COMPLETION_SCRIPT) + elif shell == "zsh": + print(_ZSH_COMPLETION_SCRIPT) + elif shell == "fish": + print(_FISH_COMPLETION_SCRIPT) + elif shell == "powershell": + print(_POWERSHELL_COMPLETION_SCRIPT) + + +def _parse_args(parser: argparse.ArgumentParser) -> argparse.Namespace: + """Parse arguments and return namespace.""" + return parser.parse_args() + + +def _main() -> None: + # Check if this is a completion source request + completion_source = _get_completion_source() + if completion_source: + args = sys.argv[1:] + if completion_source == "bash_source": + _handle_bash_source(args) + elif completion_source == "zsh_source": + _handle_zsh_source(args) + elif completion_source == "fish_source": + _handle_fish_source(args) + elif completion_source == "powershell_source": + _handle_powershell_source(args) + return + + parser = _build_parser() + parsed = _parse_args(parser) + + if hasattr(parsed, "func"): + parsed.func(parsed) + else: + parser.print_help() + + +def _get_completion_source() -> Optional[str]: + """Check if the CLI was invoked as a completion source.""" + import os + + return os.environ.get("OPENAI_COMPLETE") + + +def main() -> int: + try: + _main() + except Exception as err: + sys.stderr.write(f"Error: {err}\n") + return 1 + return 0 + + +if __name__ == "__main__": + sys.exit(main()) diff --git a/tests/test_cli.py b/tests/test_cli.py new file mode 100644 index 0000000000..d6fe97332e --- /dev/null +++ b/tests/test_cli.py @@ -0,0 +1,85 @@ +from __future__ import annotations + +import subprocess +import sys + + +def test_completion_help() -> None: + """Test that the completion subcommand shows help.""" + result = subprocess.run( + [sys.executable, "-m", "openai", "completion", "--help"], + capture_output=True, + text=True, + ) + assert result.returncode == 0 + assert "Generate shell completion scripts" in result.stdout + + +def test_completion_bash() -> None: + """Test generating bash completion script.""" + result = subprocess.run( + [sys.executable, "-m", "openai", "completion", "-s", "bash"], + capture_output=True, + text=True, + ) + assert result.returncode == 0 + assert "complete -F _openai_completion openai" in result.stdout + + +def test_completion_zsh() -> None: + """Test generating zsh completion script.""" + result = subprocess.run( + [sys.executable, "-m", "openai", "completion", "-s", "zsh"], + capture_output=True, + text=True, + ) + assert result.returncode == 0 + assert "#compdef openai" in result.stdout + + +def test_completion_fish() -> None: + """Test generating fish completion script.""" + result = subprocess.run( + [sys.executable, "-m", "openai", "completion", "-s", "fish"], + capture_output=True, + text=True, + ) + assert result.returncode == 0 + assert "complete -f -c openai" in result.stdout + + +def test_completion_powershell() -> None: + """Test generating PowerShell completion script.""" + result = subprocess.run( + [sys.executable, "-m", "openai", "completion", "-s", "powershell"], + capture_output=True, + text=True, + ) + assert result.returncode == 0 + assert "Register-ArgumentCompleter" in result.stdout + + +def test_completion_missing_shell() -> None: + """Test that completion requires --shell flag.""" + result = subprocess.run( + [sys.executable, "-m", "openai", "completion"], + capture_output=True, + text=True, + ) + assert result.returncode != 0 + + +def test_completion_bash_source() -> None: + """Test bash completion source.""" + import os + + env = os.environ.copy() + env["OPENAI_COMPLETE"] = "bash_source" + result = subprocess.run( + [sys.executable, "-m", "openai", ""], + capture_output=True, + text=True, + env=env, + ) + assert result.returncode == 0 + assert "completion" in result.stdout