44import re
55import sys
66from argparse import (
7+ ONE_OR_MORE ,
8+ REMAINDER ,
79 SUPPRESS ,
10+ ZERO_OR_MORE ,
811 Action ,
912 ArgumentParser ,
1013 _AppendAction ,
4952 _AppendConstAction ,
5053 _CountAction ,
5154)
52- OPTION_END = _HelpAction , _VersionAction
55+
56+
57+ class _ShtabPrintCompletionAction (Action ):
58+ pass
59+
60+
61+ OPTION_END = _HelpAction , _VersionAction , _ShtabPrintCompletionAction
5362OPTION_MULTI = _AppendAction , _AppendConstAction , _CountAction
5463
5564
@@ -462,11 +471,16 @@ def complete_zsh(parser, root_prefix=None, preamble="", choice_functions=None):
462471 if choice_functions :
463472 choice_type2fn .update (choice_functions )
464473
474+ def is_opt_end (opt ):
475+ return isinstance (opt , OPTION_END ) or opt .nargs == REMAINDER
476+
477+ def is_opt_multiline (opt ):
478+ return isinstance (opt , OPTION_MULTI )
479+
465480 def format_optional (opt ):
466481 return (('{nargs}{options}"[{help}]"' if isinstance (
467482 opt , FLAG_OPTION ) else '{nargs}{options}"[{help}]:{dest}:{pattern}"' ).format (
468- nargs = ('"(- :)"' if isinstance (opt , OPTION_END ) else
469- '"*"' if isinstance (opt , OPTION_MULTI ) else "" ),
483+ nargs = ('"(- : *)"' if is_opt_end (opt ) else '"*"' if is_opt_multiline (opt ) else "" ),
470484 options = ("{{{}}}" .format ("," .join (opt .option_strings ))
471485 if len (opt .option_strings ) > 1 else '"{}"' .format ("" .join (
472486 opt .option_strings ))),
@@ -480,7 +494,7 @@ def format_optional(opt):
480494
481495 def format_positional (opt ):
482496 return '"{nargs}:{help}:{pattern}"' .format (
483- nargs = {"+" : "(*)" , "*" : "(*): " }.get (opt .nargs , "" ),
497+ nargs = {ONE_OR_MORE : "(*)" , ZERO_OR_MORE : "(*):" , REMAINDER : "(-)* " }.get (opt .nargs , "" ),
484498 help = escape_zsh ((opt .help or opt .dest ).strip ().split ("\n " )[0 ]),
485499 pattern = complete2pattern (opt .complete , "zsh" , choice_type2fn ) if hasattr (
486500 opt , "complete" ) else
@@ -492,10 +506,12 @@ def format_positional(opt):
492506 all_commands = {
493507 root_prefix : {
494508 "cmd" : prog , "arguments" : [
495- format_optional (opt ) for opt in parser ._get_optional_actions ()
496- if opt .help != SUPPRESS ], "help" : (parser .description
497- or "" ).strip ().split ("\n " )[0 ], "commands" : [],
498- "paths" : []}}
509+ format_optional (opt )
510+ for opt in parser ._get_optional_actions () if opt .help != SUPPRESS ] + [
511+ format_positional (opt ) for opt in parser ._get_positional_actions ()
512+ if opt .help != SUPPRESS and opt .choices is None ],
513+ "help" : (parser .description
514+ or "" ).strip ().split ("\n " )[0 ], "commands" : [], "paths" : []}}
499515
500516 def recurse (parser , prefix , paths = None ):
501517 paths = paths or []
@@ -564,11 +580,12 @@ def command_case(prefix, options):
564580
565581 return """\
566582 {prefix}() {{
567- local context state line curcontext="$curcontext"
583+ local context state line curcontext="$curcontext" one_or_more='(-)*' remainder='(*)'
568584
569- _arguments -C ${prefix}_options \\
570- ': :{prefix}_commands' \\
571- '*::: :->{name}'
585+ if ((${{{prefix}_options[(I)${{(q)one_or_more}}*]}} + ${{{prefix}_options[(I)${{(q)remainder}}*]}} == 0)); then # noqa: E501
586+ {prefix}_options+=(': :{prefix}_commands' '*::: :->{name}')
587+ fi
588+ _arguments -C ${prefix}_options
572589
573590 case $state in
574591 {name})
@@ -766,7 +783,7 @@ def complete(parser: ArgumentParser, shell: str = "bash", root_prefix: Opt[str]
766783
767784
768785def completion_action (parent = None , preamble = "" ):
769- class PrintCompletionAction (Action ):
786+ class PrintCompletionAction (_ShtabPrintCompletionAction ):
770787 def __call__ (self , parser , namespace , values , option_string = None ):
771788 print (complete (parent or parser , values , preamble = preamble ))
772789 parser .exit (0 )
0 commit comments