Merge pull request #73 from skshetry/zsh-subcommand

Author: casperdcl

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,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
@@ -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,
     )
 
 

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", "_describe 'script.py commands' _commands",
+            "'*::: :->script.py'", "script.py)"]
     elif shell == "tcsh":
         assert script_py == ["complete script.py \\"]
     else: