diff options
Diffstat (limited to 'python/vyos/opmode.py')
| -rw-r--r-- | python/vyos/opmode.py | 40 |
1 files changed, 38 insertions, 2 deletions
diff --git a/python/vyos/opmode.py b/python/vyos/opmode.py index d02ad4de6..230a85541 100644 --- a/python/vyos/opmode.py +++ b/python/vyos/opmode.py @@ -128,6 +128,25 @@ def _get_arg_type(t): else: return t +def _is_literal_type(t): + if _is_optional_type(t): + t = _get_arg_type(t) + + if typing.get_origin(t) == typing.Literal: + return True + + return False + +def _get_literal_values(t): + """ Returns the tuple of allowed values for a Literal type + """ + if not _is_literal_type(t): + return tuple() + if _is_optional_type(t): + t = _get_arg_type(t) + + return typing.get_args(t) + def _normalize_field_name(name): # Convert the name to string if it is not # (in some cases they may be numbers) @@ -190,13 +209,30 @@ def run(module): for opt in type_hints: th = type_hints[opt] + # Function argument names use underscores as separators + # but command-line options should use hyphens + # Without this, we'd get options like "--foo_bar" + opt = re.sub(r'_', '-', opt) + if _get_arg_type(th) == bool: subparser.add_argument(f"--{opt}", action='store_true') else: if _is_optional_type(th): - subparser.add_argument(f"--{opt}", type=_get_arg_type(th), default=None) + if _is_literal_type(th): + subparser.add_argument(f"--{opt}", + choices=list(_get_literal_values(th)), + default=None) + else: + subparser.add_argument(f"--{opt}", + type=_get_arg_type(th), default=None) else: - subparser.add_argument(f"--{opt}", type=_get_arg_type(th), required=True) + if _is_literal_type(th): + subparser.add_argument(f"--{opt}", + choices=list(_get_literal_values(th)), + required=True) + else: + subparser.add_argument(f"--{opt}", + type=_get_arg_type(th), required=True) # Get options as a dict rather than a namespace, # so that we can modify it and pack for passing to functions |