commandkit/internal/cobra/testdata/completion_terminal.py

#!/usr/bin/env python3
"""Exercise one generated completion through an interactive shell."""

import os
import shlex
import shutil
import sys

import pexpect

shell, script, line = sys.argv[1:]
executable = shutil.which(shell)
if executable is None:
    raise SystemExit(f"{shell} is not installed")
quoted_script = shlex.quote(script)
if shell == "bash":
    args = ["--noprofile", "--norc", "-i"]
    bash_init = os.environ.get("COMPLETION_BASH_INIT", "/usr/share/bash-completion/bash_completion")
    if not os.path.isfile(bash_init):
        raise SystemExit("Set COMPLETION_BASH_INIT to the bash-completion initialization file")
    setup = 'PS1="PROMPT> "; source ' + shlex.quote(bash_init) + "; source " + quoted_script
elif shell == "zsh":
    args = ["-f"]
    setup = 'PROMPT="PROMPT> "; '
    if os.environ.get("COMPLETION_ZSH_MODULE_PATH"):
        setup += "module_path=(" + shlex.quote(os.environ["COMPLETION_ZSH_MODULE_PATH"]) + "); "
    if os.environ.get("COMPLETION_ZSH_FUNCTION_ROOT"):
        setup += "fpath=(" + shlex.quote(os.environ["COMPLETION_ZSH_FUNCTION_ROOT"]) + "/**/*(N-/) $fpath); "
    setup += "autoload -Uz compinit; compinit -D; source " + quoted_script
elif shell == "fish":
    args = ["--no-config", "--interactive"]
    setup = "function fish_prompt; printf '\\120\\122\\117\\115\\120\\124\\076\\040'; end; source " + quoted_script
else:
    raise SystemExit(f"unsupported shell: {shell}")

tab_only = os.environ.get("PROBE_TAB_ONLY") == "1"
if tab_only:
    if shell == "zsh":
        setup += "; function ck_tab_done { print CK_TAB_FINISHED; }; zle -N ck_tab_done; bindkey '^X^T' ck_tab_done"
    elif shell == "bash":
        setup += r"""; bind -x '"\C-x\C-t": "printf CK_TAB_FINISHED"'"""
    else:
        raise SystemExit("tab-only probe requires Bash or Zsh")

child = pexpect.spawn(executable, args, env=os.environ.copy(), cwd=os.environ["PROBE_DIR"],
                      encoding="utf-8", timeout=15, dimensions=(24, 120))
child.delaybeforesend = 0.05
try:
    child.setecho(False)
    child.sendline(setup)
    child.expect_exact("PROMPT> ")
    child.send(line)
    child.send("\t")
    if tab_only:
        child.sendcontrol("x")
        child.sendcontrol("t")
        child.expect_exact("CK_TAB_FINISHED")
    elif os.environ.get("PROBE_LIST") == "1":
        child.send("\t")
        remaining = ["Readable text", "YAML output"]
        captured = []
        while remaining:
            index = child.expect_exact(remaining)
            captured.append(child.before + child.after)
            remaining.pop(index)
        print("".join(captured))
    else:
        child.sendline("")
        child.expect_exact("PROMPT> ")
finally:
    child.close(force=True)