From 55680af7964d5881bb17531ff6584a87d8cfc087 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Saugat=20Pachhai=20=28=E0=A4=B8=E0=A5=8C=E0=A4=97=E0=A4=BE?= =?UTF-8?q?=E0=A4=A4=29?= Date: Wed, 13 Apr 2022 18:56:09 +0545 Subject: [PATCH 1/3] zsh: recurse over subcommand --- shtab/__init__.py | 208 +++++++++++++++++++++++++++----------------- tests/test_shtab.py | 4 +- 2 files changed, 132 insertions(+), 80 deletions(-) diff --git a/shtab/__init__.py b/shtab/__init__.py index 0b06359..a877e34 100644 --- a/shtab/__init__.py +++ b/shtab/__init__.py @@ -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 @@ -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: @@ -486,47 +486,124 @@ 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) + # {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 not sub.choices or not isinstance(sub.choices, dict): + # positional argument + opt = sub + if opt.help != SUPPRESS: + all_commands[prefix]["arguments"].append(format_positional(opt)) + 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) - subcommands[cmd] = { - "help": (subparser.description or "").strip().split("\n")[0], - "arguments": arguments} - log.debug("subcommands:%s:%s", cmd, subcommands[cmd]) + preamble = """\ +# Custom Preamble - log.debug("subcommands:%s:%s", root_prefix, sorted(subcommands)) +{} +# 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 @@ -538,49 +615,22 @@ def format_positional(opt): # AUTOMATCALLY GENERATED by `shtab` -${root_prefix}_options_=( - ${root_options} -) +${command_cases} -${root_prefix}_commands_() { - local _commands=( - ${commands} - ) +${command_commands} - _describe '${prog} commands' _commands -} -${subcommands} -${preamble} -typeset -A opt_args -local context state line curcontext="$curcontext" +${command_options} -_arguments \\ - $$${root_prefix}_options_ \\ - ${root_arguments} \\ - ': :${root_prefix}_commands_' \\ - '*::args:->args' +${preamble} -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, ) diff --git a/tests/test_shtab.py b/tests/test_shtab.py index 721fa77..87f94e2 100644 --- a/tests/test_shtab.py +++ b/tests/test_shtab.py @@ -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", "'*::: :->script.py'", "script.py)", + "_describe 'script.py commands' _commands"] elif shell == "tcsh": assert script_py == ["complete script.py \\"] else: From 0e5d2239ef8939031c6269fd8556aa323131e683 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Saugat=20Pachhai=20=28=E0=A4=B8=E0=A5=8C=E0=A4=97=E0=A4=BE?= =?UTF-8?q?=E0=A4=A4=29?= Date: Wed, 13 Apr 2022 19:25:20 +0545 Subject: [PATCH 2/3] avoid duplicating positional help message --- shtab/__init__.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/shtab/__init__.py b/shtab/__init__.py index a877e34..600b5b5 100644 --- a/shtab/__init__.py +++ b/shtab/__init__.py @@ -499,11 +499,11 @@ 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 - opt = sub - if opt.help != SUPPRESS: - all_commands[prefix]["arguments"].append(format_positional(opt)) + all_commands[prefix]["arguments"].append(format_positional(sub)) else: # subparser log.debug("choices:{}:{}".format(prefix, sorted(sub.choices))) public_cmds = get_public_subcommands(sub) From 6a975065260a443afb2c77dc26e91f6a2f18e92a Mon Sep 17 00:00:00 2001 From: Casper da Costa-Luis Date: Tue, 19 Apr 2022 12:50:30 +0100 Subject: [PATCH 3/3] minor reordering --- shtab/__init__.py | 4 +--- tests/test_shtab.py | 4 ++-- 2 files changed, 3 insertions(+), 5 deletions(-) diff --git a/shtab/__init__.py b/shtab/__init__.py index 600b5b5..5467cfd 100644 --- a/shtab/__init__.py +++ b/shtab/__init__.py @@ -599,7 +599,6 @@ def command_list(prefix, options): preamble = """\ # Custom Preamble - {} # End Custom Preamble @@ -615,12 +614,11 @@ def command_list(prefix, options): # AUTOMATCALLY GENERATED by `shtab` -${command_cases} - ${command_commands} ${command_options} +${command_cases} ${preamble} typeset -A opt_args diff --git a/tests/test_shtab.py b/tests/test_shtab.py index 87f94e2..fc8d09c 100644 --- a/tests/test_shtab.py +++ b/tests/test_shtab.py @@ -91,8 +91,8 @@ def test_prog_scripts(shell, caplog, capsys): assert script_py == ["complete -o filenames -F _shtab_shtab script.py"] elif shell == "zsh": assert script_py == [ - "#compdef script.py", "'*::: :->script.py'", "script.py)", - "_describe 'script.py commands' _commands"] + "#compdef script.py", "_describe 'script.py commands' _commands", + "'*::: :->script.py'", "script.py)"] elif shell == "tcsh": assert script_py == ["complete script.py \\"] else: