Skip to content

Commit

Permalink
Merge pull request #73 from skshetry/zsh-subcommand
Browse files Browse the repository at this point in the history
  • Loading branch information
@casperdcl
casperdcl authored Apr 19, 2022
2 parents c7a390a + 6a97506 commit cf19980
Show file tree
Hide file tree
Showing 2 changed files with 131 additions and 81 deletions.
208 changes: 128 additions & 80 deletions shtab/__init__.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -16,6 +16,7 @@
)
from collections import defaultdict
from functools import total_ordering
from itertools import starmap
from string import Template
from typing import Any, Dict, List
from typing import Optional as Opt
Expand Down Expand Up @@ -452,9 +453,8 @@ def complete_zsh(parser, root_prefix=None, preamble="", choice_functions=None):
See `complete` for arguments.
"""
root_prefix = wordify("_shtab_" + (root_prefix or parser.prog))
root_arguments = []
subcommands = {} # {cmd: {"help": help, "arguments": [arguments]}}
prog = parser.prog
root_prefix = wordify("_shtab_" + (root_prefix or prog))

choice_type2fn = {k: v["zsh"] for k, v in CHOICE_FUNCTIONS.items()}
if choice_functions:
Expand Down Expand Up @@ -486,47 +486,123 @@ def format_positional(opt):
"({})".format(" ".join(map(str, opt.choices)))) if opt.choices else "",
)

for sub in parser._get_positional_actions():
if not sub.choices or not isinstance(sub.choices, dict):
# positional argument
opt = sub
if opt.help != SUPPRESS:
root_arguments.append(format_positional(opt))
else: # subparser
log.debug("choices:{}:{}".format(root_prefix, sorted(sub.choices)))
public_cmds = get_public_subcommands(sub)
for cmd, subparser in sub.choices.items():
if cmd not in public_cmds:
log.debug("skip:subcommand:%s", cmd)
continue
log.debug("subcommand:%s", cmd)

# optionals
arguments = [
format_optional(opt) for opt in subparser._get_optional_actions()
if opt.help != SUPPRESS]

# subcommand positionals
subsubs = sum(
(list(opt.choices) for opt in subparser._get_positional_actions()
if isinstance(opt.choices, dict)),
[],
)
if subsubs:
arguments.append('"1:Sub command:({})"'.format(" ".join(subsubs)))

# positionals
arguments.extend(
format_positional(opt) for opt in subparser._get_positional_actions()
if not isinstance(opt.choices, dict) if opt.help != SUPPRESS)

subcommands[cmd] = {
"help": (subparser.description or "").strip().split("\n")[0],
"arguments": arguments}
log.debug("subcommands:%s:%s", cmd, subcommands[cmd])
# {cmd: {"help": help, "arguments": [arguments]}}
all_commands = {
root_prefix: {
"cmd": prog, "arguments": [
format_optional(opt) for opt in parser._get_optional_actions()
if opt.help != SUPPRESS], "help": (parser.description
or "").strip().split("\n")[0], "commands": [],
"paths": []}}

def recurse(parser, prefix, paths=None):
paths = paths or []
subcmds = []
for sub in parser._get_positional_actions():
if sub.help == SUPPRESS or not sub.choices:
continue
if not sub.choices or not isinstance(sub.choices, dict):
# positional argument
all_commands[prefix]["arguments"].append(format_positional(sub))
else: # subparser
log.debug("choices:{}:{}".format(prefix, sorted(sub.choices)))
public_cmds = get_public_subcommands(sub)
for cmd, subparser in sub.choices.items():
if cmd not in public_cmds:
log.debug("skip:subcommand:%s", cmd)
continue
log.debug("subcommand:%s", cmd)

# optionals
arguments = [
format_optional(opt) for opt in subparser._get_optional_actions()
if opt.help != SUPPRESS]

# positionals
arguments.extend(
format_positional(opt) for opt in subparser._get_positional_actions()
if not isinstance(opt.choices, dict) if opt.help != SUPPRESS)

new_pref = prefix + "_" + wordify(cmd)
options = all_commands[new_pref] = {
"cmd": cmd, "help": (subparser.description or "").strip().split("\n")[0],
"arguments": arguments, "paths": [*paths, cmd]}
new_subcmds = recurse(subparser, new_pref, [*paths, cmd])
options["commands"] = {
all_commands[pref]["cmd"]: all_commands[pref]
for pref in new_subcmds if pref in all_commands}
subcmds.extend([*new_subcmds, new_pref])
log.debug("subcommands:%s:%s", cmd, options)
return subcmds

recurse(parser, root_prefix)
all_commands[root_prefix]["commands"] = {
options["cmd"]: options
for prefix, options in sorted(all_commands.items())
if len(options.get("paths", [])) < 2 and prefix != root_prefix}
subcommands = {
prefix: options
for prefix, options in all_commands.items() if options.get("commands")}
subcommands.setdefault(root_prefix, all_commands[root_prefix])
log.debug("subcommands:%s:%s", root_prefix, sorted(all_commands))

def command_case(prefix, options):
name = options["cmd"]
commands = options["commands"]
case_fmt_on_no_sub = """{name}) _arguments -C ${prefix}_{name}_options ;;"""
case_fmt_on_sub = """{name}) {prefix}_{name} ;;"""

cases = []
for _, options in sorted(commands.items()):
fmt = case_fmt_on_sub if options.get("commands") else case_fmt_on_no_sub
cases.append(fmt.format(name=options["cmd"], prefix=prefix))
cases = "\n\t".expandtabs(8).join(cases)

return """\
{prefix}() {{
local context state line curcontext="$curcontext"
_arguments -C ${prefix}_options \\
': :{prefix}_commands' \\
'*::: :->{name}'
case $state in
{name})
words=($line[1] "${{words[@]}}")
(( CURRENT += 1 ))
curcontext="${{curcontext%:*:*}}:{prefix}-$line[1]:"
case $line[1] in
{cases}
esac
esac
}}
""".format(prefix=prefix, name=name, cases=cases)

def command_option(prefix, options):
return """\
{prefix}_options=(
{arguments}
)
""".format(prefix=prefix, arguments="\n ".join(options["arguments"]))

def command_list(prefix, options):
name = " ".join([prog, *options["paths"]])
commands = "\n ".join('"{}:{}"'.format(cmd, escape_zsh(opt["help"]))
for cmd, opt in sorted(options["commands"].items()))
return """
{prefix}_commands() {{
local _commands=(
{commands}
)
_describe '{name} commands' _commands
}}""".format(prefix=prefix, name=name, commands=commands)

log.debug("subcommands:%s:%s", root_prefix, sorted(subcommands))
preamble = """\
# Custom Preamble
{}
# End Custom Preamble
""".format(preamble.rstrip()) if preamble else ""
# References:
# - https://github.com/zsh-users/zsh-completions
# - http://zsh.sourceforge.net/Doc/Release/Completion-System.html
Expand All @@ -538,49 +614,21 @@ def format_positional(opt):
# AUTOMATCALLY GENERATED by `shtab`
${root_prefix}_options_=(
${root_options}
)
${command_commands}
${root_prefix}_commands_() {
local _commands=(
${commands}
)
${command_options}
_describe '${prog} commands' _commands
}
${subcommands}
${command_cases}
${preamble}
typeset -A opt_args
local context state line curcontext="$curcontext"
_arguments \\
$$${root_prefix}_options_ \\
${root_arguments} \\
': :${root_prefix}_commands_' \\
'*::args:->args'
case $words[1] in
${commands_case}
esac""").safe_substitute(
typeset -A opt_args
${root_prefix} "$@\"""").safe_substitute(
prog=prog,
root_prefix=root_prefix,
prog=parser.prog,
commands="\n ".join('"{}:{}"'.format(cmd, escape_zsh(subcommands[cmd]["help"]))
for cmd in sorted(subcommands)),
root_arguments=" \\\n ".join(root_arguments),
root_options="\n ".join(
format_optional(opt) for opt in parser._get_optional_actions()
if opt.help != SUPPRESS),
commands_case="\n ".join("{cmd_orig}) _arguments ${root_prefix}_{cmd} ;;".format(
cmd_orig=cmd, cmd=wordify(cmd), root_prefix=root_prefix)
for cmd in sorted(subcommands)),
subcommands="\n".join("""
{root_prefix}_{cmd}=(
{arguments}
)""".format(root_prefix=root_prefix, cmd=wordify(cmd), arguments="\n ".join(
subcommands[cmd]["arguments"])) for cmd in sorted(subcommands)),
preamble=("\n# Custom Preamble\n" + preamble +
"\n# End Custom Preamble\n" if preamble else ""),
command_cases="\n".join(starmap(command_case, sorted(subcommands.items()))),
command_commands="\n".join(starmap(command_list, sorted(subcommands.items()))),
command_options="\n".join(starmap(command_option, sorted(all_commands.items()))),
preamble=preamble,
)


Expand Down
4 changes: 3 additions & 1 deletion tests/test_shtab.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -90,7 +90,9 @@ def test_prog_scripts(shell, caplog, capsys):
if shell == "bash":
assert script_py == ["complete -o filenames -F _shtab_shtab script.py"]
elif shell == "zsh":
assert script_py == ["#compdef script.py", "_describe 'script.py commands' _commands"]
assert script_py == [
"#compdef script.py", "_describe 'script.py commands' _commands",
"'*::: :->script.py'", "script.py)"]
elif shell == "tcsh":
assert script_py == ["complete script.py \\"]
else:
Expand Down

0 comments on commit cf19980

Please sign in to comment.