shtab completions improved for unions with open values and literals (#851)

Author: mauvilsa

CHANGELOG.rst

@@ -15,6 +15,13 @@ paths are considered internals and can change in minor and patch releases.
 v4.47.0 (unreleased)
 --------------------
 
+Added
+^^^^^
+- Improved ``shtab`` bash typehint completions for ``Union`` and ``Literal``:
+  choices for unions that also accept open values now require a prefix to be
+  completed, and literal booleans/``None`` now complete as ``false``/``true``
+  and ``null`` (`#851 <https://github.com/omni-us/jsonargparse/pull/851>`__).
+
 Fixed
 ^^^^^
 - Positional arguments with ``nargs="*"`` or ``nargs="?"`` now correctly allow

jsonargparse/_completions.py

@@ -174,12 +174,18 @@ def shtab_prepare_action(action, parser) -> None:
     if isinstance(action, ActionTypeHint):
         skip = getattr(action, "sub_add_kwargs", {}).get("skip", set())
         prefix = action.option_strings[0] if action.option_strings else None
-        choices = get_typehint_choices(action._typehint, prefix, parser, skip)
+        choices, require_prefix = get_typehint_choices(action._typehint, prefix, parser, skip)
         if shtab_shell.get() == "bash":
             message = f"Expected type: {type_to_str(action._typehint)}"
             if action.option_strings == []:
                 message = f"Argument: {action.dest}; " + message
-            add_bash_typehint_completion(parser, action, message, choices)
+            add_bash_typehint_completion(
+                parser,
+                action,
+                message,
+                choices,
+                require_prefix=require_prefix,
+            )
             choices = None
     elif isinstance(action, _ActionHelpClassPath):
         choices = get_help_class_choices(action._typehint)
@@ -197,15 +203,23 @@ def shtab_prepare_action(action, parser) -> None:
   fi
 }
 %(name)s() {
-  local MATCH=( $(IFS=" " compgen -W "$1" "$2") )
+  local REQUIRE_PREFIX="$4"
+  local MATCH=()
+  if [ "$REQUIRE_PREFIX" = 1 ] && [ -z "$2" ]; then
+    MATCH=()
+  else
+    MATCH=( $(IFS=" " compgen -W "$1" "$2") )
+  fi
   if [ ${#MATCH[@]} = 0 ]; then
     if [ "$COMP_TYPE" = 63 ]; then
       MATCHED=$(_jsonargparse_%%s_matched_choices "$1" "${MATCH[*]}")
       printf "%(b)s\\n$3$MATCHED\\n%(n)s" >&2
       kill -WINCH $$
     fi
   else
-    IFS=" " compgen -W "$1" "$2"
+    for match in "${MATCH[@]}"; do
+      echo "$match"
+    done
     if [ "$COMP_TYPE" = 63 ]; then
       MATCHED=$(_jsonargparse_%%s_matched_choices "$1" "${MATCH[*]}")
       printf "%(b)s\\n$3$MATCHED%(n)s" >&2
@@ -219,53 +233,77 @@ def shtab_prepare_action(action, parser) -> None:
 }
 
 
-def add_bash_typehint_completion(parser, action, message, choices) -> None:
+def add_bash_typehint_completion(parser, action, message, choices, require_prefix=False) -> None:
     fn_typehint = norm_name(bash_compgen_typehint_name % shtab_prog.get())
     fn_name = parser.prog.replace(" [options] ", "_")
     fn_name = norm_name(f"_jsonargparse_{fn_name}_{action.dest}_typehint")
-    fn = '{fn_name}(){{ {fn_typehint} "{choices}" "$1" "{message}"; }}'.format(
+    fn = '{fn_name}(){{ {fn_typehint} "{choices}" "$1" "{message}" {require_prefix}; }}'.format(
         fn_name=fn_name,
         fn_typehint=fn_typehint,
         choices=" ".join(choices),
         message=message,
+        require_prefix=1 if require_prefix else 0,
     )
     shtab_preambles.get().append(fn)
     action.complete = {"bash": fn_name}
 
 
-def get_typehint_choices(typehint, prefix, parser, skip, choices=None, added_subclasses=None) -> list[str]:
-    if choices is None:
-        choices = []
+def get_typehint_choices(typehint, prefix, parser, skip, added_subclasses=None) -> tuple[list[str], bool]:
     if not added_subclasses:
         added_subclasses = set()
-    if typehint is bool:
-        choices.extend(["true", "false"])
-    elif typehint is NoneType:
-        choices.append("null")
-    elif is_subclass(typehint, Enum):
-        choices.extend(list(typehint.__members__))
-    else:
+
+    def get_choices_state(typehint) -> tuple[list[str], bool, bool]:
+        if typehint is bool:
+            return ["true", "false"], True, False
+        if typehint is NoneType:
+            return ["null"], True, False
+        if is_subclass(typehint, Enum):
+            return list(typehint.__members__), True, False
+
         origin = get_typehint_origin(typehint)
         if origin == Literal:
-            choices.extend([str(a) for a in typehint.__args__ if isinstance(a, (str, int, float))])
-        elif origin == Union:
+            choices = []
+            for arg in typehint.__args__:
+                if isinstance(arg, bool):
+                    choices.append(str(arg).lower())
+                elif arg is None:
+                    choices.append("null")
+                elif isinstance(arg, (str, int, float)):
+                    choices.append(str(arg))
+            return choices, True, False
+
+        if origin == Union:
+            choices = []
+            has_explicit_choices = False
+            has_open_values = False
             for subtype in typehint.__args__:
                 if subtype in added_subclasses or subtype is object:
                     continue
-                get_typehint_choices(subtype, prefix, parser, skip, choices, added_subclasses)
-        elif ActionTypeHint.is_subclass_typehint(typehint):
+                subchoices, subexplicit, subopen = get_choices_state(subtype)
+                choices.extend(subchoices)
+                has_explicit_choices = has_explicit_choices or subexplicit
+                has_open_values = has_open_values or subopen
+            return choices, has_explicit_choices, has_open_values
+
+        if ActionTypeHint.is_subclass_typehint(typehint):
             added_subclasses.add(typehint)
-            choices.extend(add_subactions_and_get_subclass_choices(typehint, prefix, parser, skip, added_subclasses))
-        elif origin in callable_origin_types:
+            choices = add_subactions_and_get_subclass_choices(typehint, prefix, parser, skip, added_subclasses)
+            return choices, True, False
+
+        if origin in callable_origin_types:
             return_type = get_callable_return_type(typehint)
             if return_type and ActionTypeHint.is_subclass_typehint(return_type):
                 num_args = len(typehint.__args__) - 1
                 skip.add(num_args)
-                choices.extend(
-                    add_subactions_and_get_subclass_choices(return_type, prefix, parser, skip, added_subclasses)
-                )
+                choices = add_subactions_and_get_subclass_choices(return_type, prefix, parser, skip, added_subclasses)
+                return choices, True, False
+            return [], False, return_type is None
+
+        return [], False, True
 
-    return [] if choices == ["null"] else choices
+    choices, has_explicit_choices, has_open_values = get_choices_state(typehint)
+    require_prefix = get_typehint_origin(typehint) == Union and has_explicit_choices and has_open_values
+    return choices, require_prefix
 
 
 def add_subactions_and_get_subclass_choices(typehint, prefix, parser, skip, added_subclasses) -> list[str]:
@@ -295,11 +333,19 @@ def add_subactions_and_get_subclass_choices(typehint, prefix, parser, skip, adde
             if option_string not in parser._option_string_actions:
                 action = parser.add_argument(option_string)
                 for subtype in unique(subtypes):
-                    subchoices = get_typehint_choices(subtype, option_string, parser, skip, None, added_subclasses)
+                    subchoices, require_prefix = get_typehint_choices(
+                        subtype, option_string, parser, skip, added_subclasses
+                    )
                     if shtab_shell.get() == "bash":
                         message = f"Expected type: {type_to_str(subtype)}; "
                         message += f"Accepted by subclasses: {', '.join(subclasses[name])}"
-                        add_bash_typehint_completion(parser, action, message, subchoices)
+                        add_bash_typehint_completion(
+                            parser,
+                            action,
+                            message,
+                            subchoices,
+                            require_prefix=require_prefix,
+                        )
                     elif subchoices:
                         action.choices = subchoices
 

jsonargparse_tests/test_shtab.py

@@ -115,6 +115,18 @@ def test_bash_optional_bool(parser, subtests):
     )
 
 
+def test_bash_optional_int(parser, subtests):
+    parser.add_argument("--num", type=Optional[int])
+    assert_bash_typehint_completions(
+        subtests,
+        parser,
+        [
+            ("num", Optional[int], "", [], "0/1"),
+            ("num", Optional[int], "n", ["null"], "1/1"),
+        ],
+    )
+
+
 def test_bash_argument_group(parser, subtests):
     group = parser.add_argument_group("Group1")
     group.add_argument("--bool", type=bool)
@@ -170,6 +182,19 @@ def test_bash_literal(parser, subtests):
     )
 
 
+def test_bash_literal_none(parser, subtests):
+    typehint = Literal[None]
+    parser.add_argument("--literal", type=typehint)
+    assert_bash_typehint_completions(
+        subtests,
+        parser,
+        [
+            ("literal", typehint, "", ["null"], "1/1"),
+            ("literal", typehint, "n", ["null"], "1/1"),
+        ],
+    )
+
+
 def test_bash_union(parser, subtests):
     typehint = Optional[Union[bool, AXEnum]]
     parser.add_argument("--union", type=typehint)
@@ -183,6 +208,32 @@ def test_bash_union(parser, subtests):
     )
 
 
+def test_bash_union_literal_and_int(parser, subtests):
+    typehint = Union[Literal[False], int]
+    parser.add_argument("--union", type=typehint)
+    assert_bash_typehint_completions(
+        subtests,
+        parser,
+        [
+            ("union", typehint, "", [], "0/1"),
+            ("union", typehint, "f", ["false"], "1/1"),
+        ],
+    )
+
+
+def test_bash_union_float_and_enum(parser, subtests):
+    typehint = Union[float, AXEnum]
+    parser.add_argument("--union", type=typehint)
+    assert_bash_typehint_completions(
+        subtests,
+        parser,
+        [
+            ("union", typehint, "", [], "0/3"),
+            ("union", typehint, "X", ["XY", "XZ"], "2/3"),
+        ],
+    )
+
+
 def test_bash_positional(parser, subtests):
     typehint = Literal["Alice", "Bob"]
     parser.add_argument("name", type=typehint)
@@ -363,6 +414,19 @@ def test_bash_callable_return_class(parser, subtests):
     )
 
 
+def test_bash_callable_return_int(parser, subtests):
+    typehint = Callable[[int], int]
+    parser.add_argument("--num", type=typehint)
+    assert_bash_typehint_completions(
+        subtests,
+        parser,
+        [
+            ("num", typehint, "", [], None),
+            ("num", typehint, "1", [], None),
+        ],
+    )
+
+
 def test_bash_subcommands(parser, subparser, subtests):
     subparser.add_argument("--enum", type=AXEnum)
     subparser2 = ArgumentParser()