"""Rules for taking high-level field definitions and lowering them into inputs for argparse's `add_argument()`.""" from __future__ import annotations import collections.abc import dataclasses import json import shlex from functools import cached_property from typing import ( Any, Callable, Iterable, Literal, Optional, Sequence, Tuple, TypeVar, Union, cast, ) from typing_extensions import get_args, get_origin from . import _fields, _settings, _singleton, _strings from . import _fmtlib as fmt from ._backends import _argparse as argparse from ._typing_compat import is_typing_union from .conf import _markers from .constructors import ( ConstructorRegistry, PrimitiveTypeInfo, UnsupportedTypeAnnotationError, ) _T = TypeVar("_T") def flag_to_inverse(option_string: str) -> str | None: """Converts --flag to --no-flag, --child.flag to --child.no-flag, etc. Returns None for short flags (those not starting with '--'), since single-letter aliases like -f cannot have a meaningful inverse form. """ # Short flags (like -f) cannot be inverted. if not option_string.startswith("--"): return None if "." not in option_string: option_string = "--no" + _strings.get_delimiter() + option_string[2:] else: # Loose heuristic for where to add the no-/no_ prefix. left, _, right = option_string.rpartition(".") option_string = left + ".no" + _strings.get_delimiter() + right return option_string class BooleanOptionalAction(argparse.Action): """Adapted from https://github.com/python/cpython/pull/27672""" def __init__( self, option_strings: Sequence[str], dest: str, default: _T | str | None = None, type: Callable[[str], _T] | argparse.FileType | None = None, choices: Iterable[_T] | None = None, required: bool = False, help: str | None = None, metavar: str | tuple[str, ...] | None = None, ) -> None: _option_strings = [] self._no_strings = set() for option_string in option_strings: _option_strings.append(option_string) if option_string.startswith("--"): option_string = flag_to_inverse(option_string) self._no_strings.add(option_string) _option_strings.append(option_string) super().__init__( option_strings=_option_strings, dest=dest, nargs=0, default=default, type=type, choices=choices, required=required, help=help, metavar=metavar, ) def __call__(self, parser, namespace, values, option_string=None): if option_string in self.option_strings: assert option_string is not None setattr(namespace, self.dest, option_string not in self._no_strings) @dataclasses.dataclass(frozen=True) class ArgumentDefinition: """Structure containing everything needed to define an argument.""" intern_prefix: str # True prefix. (eg for the argument's dest field) extern_prefix: str # User-facing prefix. subcommand_prefix: str # Prefix for nesting. field: _fields.FieldDefinition def get_output_key(self) -> str: """Get key used for this arg in the parsed output dict.""" if self.is_positional(): return self.lowered.name_or_flags[-1] else: assert self.lowered.dest is not None return self.lowered.dest def is_positional(self) -> bool: """Returns True if the argument should be positional in the commandline.""" return ( # Explicit positionals. _markers.Positional in self.field.markers or ( # Make required arguments positional. _markers.PositionalRequiredArgs in self.field.markers and _singleton.is_missing(self.field.default) ) # Arguments with no names (like in DummyWrapper) should be # positional. or ( self.field.extern_name == "" and self.field.intern_name == "__tyro_dummy_inner__" ) ) def add_argument( self, parser: Union[argparse.ArgumentParser, argparse._ArgumentGroup] ) -> None: """Add a defined argument to a parser.""" # Get keyword arguments, with None values removed. kwargs = dict(self.lowered.__dict__) # type: ignore kwargs.pop("instance_from_str") kwargs.pop("str_from_instance") kwargs = {k: v for k, v in kwargs.items() if v is not None} name_or_flags = kwargs.pop("name_or_flags") # Evaluate lazy help if callable. if "help" in kwargs and callable(kwargs["help"]): kwargs["help"] = kwargs["help"]() if self.is_positional(): if "required" in kwargs: kwargs.pop("required") # Can't be passed in for positional arguments. if len(name_or_flags) > 1: import warnings warnings.warn( f"Aliases were specified, but {name_or_flags} is positional. Aliases will be ignored." ) name_or_flags = name_or_flags[-1:] if kwargs.get("action", None) == "boolean_optional_action": kwargs["action"] = BooleanOptionalAction # Add argument, with aliases if available. arg = parser.add_argument(*name_or_flags, **kwargs) # Do our best to tab complete paths. # There will be false positives here, but if choices is unset they should be # harmless. # Note: shtab is now optional, so we only set completion hints if available. if "choices" not in kwargs: name_suggests_dir = ( # The conditions are intended to be conservative; if a directory path is # registered as a normal file one that's OK, the reverse on the other # hand will be overly restrictive. self.field.intern_name.endswith("_dir") or self.field.intern_name.endswith("_directory") or self.field.intern_name.endswith("_folder") ) name_suggests_path = ( self.field.intern_name.endswith("_file") or self.field.intern_name.endswith("_path") or self.field.intern_name.endswith("_filename") or name_suggests_dir ) complete_as_path = ( # Catch types like Path, List[Path], Tuple[Path, ...] etc. "Path" in str(self.field.type_stripped) # For string types, we require more evidence. or ("str" in str(self.field.type_stripped) and name_suggests_path) ) if complete_as_path: try: import shtab arg.complete = shtab.DIRECTORY if name_suggests_dir else shtab.FILE # type: ignore except ImportError: # shtab is optional; if not available, skip completion hints. pass @cached_property def lowered(self) -> LoweredArgumentDefinition: """Lowered argument definition, generated by applying a sequence of rules.""" # Each rule will mutate the lowered object. This is (unfortunately) # much faster than a functional approach. lowered = LoweredArgumentDefinition() _rule_handle_boolean_flags(self, lowered) _rule_apply_primitive_specs(self, lowered) _rule_counters(self, lowered) _rule_generate_helptext(self, lowered) _rule_set_name_or_flag_and_dest(self, lowered) _rule_positional_special_handling(self, lowered) _rule_apply_argconf(self, lowered) return lowered def is_suppressed(self) -> bool: """Returns if the argument is suppressed. Suppressed arguments won't be added to the parser.""" return _markers.Suppress in self.field.markers or ( _markers.SuppressFixed in self.field.markers and self.lowered.is_fixed() ) def get_invocation_text(self) -> tuple[fmt._Text, fmt._Text]: """Returns (invocation short, invocation long).""" # Use bright red for fixed arguments, otherwise just bold. metavar_style = ("bright_red", "bold") if self.lowered.is_fixed() else ("bold",) if self.is_positional(): assert self.lowered.metavar is not None invocation_short = fmt.text[metavar_style](self.lowered.metavar) invocation_long = fmt.text(self.lowered.metavar) return invocation_short, invocation_long name_or_flags: list[str] = list(self.lowered.name_or_flags) if self.lowered.action == "boolean_optional_action": name_or_flags = [] for name_or_flag in self.lowered.name_or_flags: name_or_flags.append(name_or_flag) # Short flags (like -f) cannot be inverted. inv = flag_to_inverse(name_or_flag) if inv is not None: name_or_flags.append(inv) # Find the first invertible flag for the short invocation display. first_inv = None for name_or_flag in self.lowered.name_or_flags: first_inv = flag_to_inverse(name_or_flag) if first_inv is not None: break # There should always be at least one long flag (starting with --) because: # - The main argument name always gets -- prepended (see _rule_generate_name_or_flag) # - tyro.conf.arg(name=...) prepends -- to the custom name # - tyro.conf.arg(aliases=...) only adds flags, never removes the main one assert first_inv is not None invocation_short = fmt.text( self.lowered.name_or_flags[0], " | ", first_inv, ) elif self.lowered.metavar is not None: invocation_short = fmt.text( self.lowered.name_or_flags[0], " ", fmt.text[metavar_style](self.lowered.metavar), ) else: invocation_short = fmt.text(self.lowered.name_or_flags[0]) if self.lowered.required is not True: invocation_short = fmt.text("[", invocation_short, "]") invocation_long_parts: list[str | fmt._Text] = [] for i, name in enumerate(name_or_flags): if i > 0: invocation_long_parts.append(", ") invocation_long_parts.append(name) if self.lowered.metavar is not None and self.lowered.metavar != "": invocation_long_parts.append(" ") invocation_long_parts.append( fmt.text[metavar_style](self.lowered.metavar) ) return invocation_short, fmt.text(*invocation_long_parts) @dataclasses.dataclass class LoweredArgumentDefinition: """Contains fields meant to be passed directly into argparse.""" # Action that is called on parsed arguments. This handles conversions from strings # to our desired types. # # The main reason we use this instead of the standard 'type' argument is to enable # mixed-type tuples. instance_from_str: Optional[Callable] = None str_from_instance: Optional[Callable] = None def is_fixed(self) -> bool: """If the instantiator is set to `None`, even after all argument transformations, it means that we don't have a valid instantiator for an argument. We then mark the argument as 'fixed', with a value always equal to the field default.""" return self.instance_from_str is None # From here on out, all fields correspond 1:1 to inputs to argparse's # add_argument() method. name_or_flags: Tuple[str, ...] = () default: Optional[Any] = None dest: Optional[str] = None required: Optional[bool] = None action: Optional[ Literal[ "count", "append", "store_true", "store_false", "boolean_optional_action" ] ] = None nargs: Optional[Union[int, Literal["*", "?"]]] = None choices: Optional[Tuple[str, ...]] = None # Note: unlike in vanilla argparse, our metavar is always a string. We handle # sequences, multiple arguments, etc, manually. metavar: Optional[str] = None help: Optional[str] | Callable[[], str] = None def _get_single_non_none_union_arg(typ: Any) -> Optional[Any]: """If typ is Union[T, None], return T. Otherwise return None.""" if not is_typing_union(get_origin(typ)): return None type_args = get_args(typ) non_none_args = [t for t in type_args if t is not type(None)] if len(non_none_args) == 1: return non_none_args[0] return None def _rule_handle_boolean_flags( arg: ArgumentDefinition, lowered: LoweredArgumentDefinition, ) -> None: # Check if this is a plain bool type. is_bool = arg.field.type_stripped is bool # Check if this is DisallowNone[bool | None], which should behave like bool # from a CLI perspective since the user can only choose True or False. if not is_bool and _markers.DisallowNone in arg.field.markers: inner_type = _get_single_non_none_union_arg(arg.field.type_stripped) is_bool = inner_type is bool if not is_bool: return if ( _singleton.is_missing(arg.field.default) or arg.is_positional() or _markers.FlagConversionOff in arg.field.markers or _markers.Fixed in arg.field.markers ): # Treat bools as a normal parameter. return # Default `False` => --flag passed in flips to `True`. if _markers.FlagCreatePairsOff in arg.field.markers: # If default is True, --flag will flip to `False`. # If default is False, --no-flag will flip to `True`. lowered.action = "store_false" if arg.field.default else "store_true" else: # Create both --flag and --no-flag. lowered.action = "boolean_optional_action" lowered.instance_from_str = lambda x: x # argparse will directly give us a bool! lowered.default = arg.field.default return def _rule_apply_primitive_specs( arg: ArgumentDefinition, lowered: LoweredArgumentDefinition, ) -> None: """The bulkiest bit: recursively analyze the type annotation and use it to determine how to instantiate it given some string from the commandline. Important: as far as argparse is concerned, all inputs are strings. Conversions from strings to our desired types happen in the instantiator; this is a bit more flexible, and lets us handle more complex types like enums and multi-type tuples.""" if _markers.Fixed in arg.field.markers: lowered.instance_from_str = None lowered.metavar = "{fixed}" lowered.required = False lowered.default = _singleton.MISSING return if lowered.instance_from_str is not None: lowered.required = False return # For positional arguments, UseAppendAction should be ignored since you can't # repeat a positional like `--flag val1 --flag val2`. Generate spec without it. exclude_markers = ( {_markers.UseAppendAction} if arg.is_positional() and _markers.UseAppendAction in arg.field.markers else None ) spec = ConstructorRegistry.get_primitive_spec( PrimitiveTypeInfo.make( cast(type, arg.field.type), arg.field.markers, exclude_markers=exclude_markers, ) ) if isinstance(spec, UnsupportedTypeAnnotationError): error = spec if _singleton.is_missing(arg.field.default): field_name = _strings.make_field_name( [arg.extern_prefix, arg.field.extern_name] ) if field_name != "": raise UnsupportedTypeAnnotationError( ( fmt.text( "Unsupported type annotation for field ", fmt.text["magenta", "bold"](field_name), " with type ", fmt.text["cyan"](str(arg.field.type)), ), *error.message, fmt.text( "To suppress this error, assign the field either a default value or a different type" ), ) ) else: # If the field name is empty, it means we're raising an error # for the direct input to `tyro.cli()`. We don't need to write # out which specific field we're complaining about. raise UnsupportedTypeAnnotationError( ( *error.message, fmt.text( "To suppress this error, assign the field either a default value or a different type" ), ) ) else: # For fields with a default, we'll get by even if there's no instantiator # available. lowered.metavar = "{fixed}" lowered.required = False lowered.default = _singleton.MISSING return # Mark lowered as required if a default is missing. if ( _singleton.is_missing(arg.field.default) and _markers._OPTIONAL_GROUP not in arg.field.markers ): lowered.required = True # We're actually going to skip the default field: if an argument is unset, the # MISSING value will be detected in _calling.py and the field default will # directly be used. This helps reduce the likelihood of issues with converting # the field default to a string format, then back to the desired type. if spec._action == "append": lowered.default = [] else: lowered.default = _singleton.MISSING_NONPROP if spec._action == "append": def append_instantiator(x: list[list[str]]) -> Any: """Handle UseAppendAction effects.""" # We'll assume that the type is annotated as Dict[...], Tuple[...], List[...], etc. container_type = get_origin(arg.field.type_stripped) if container_type is None: # Raw annotation, like `UseAppendAction[list]`. It's unlikely # that a user would use this but we can handle it. container_type = arg.field.type_stripped # Instantiate initial output. out = ( arg.field.default if not _singleton.is_missing(arg.field.default) else None ) if out is None: out = {} if container_type is dict else [] elif isinstance(out, dict): out = out.copy() else: # All sequence types will be lists for now to make sure we can # append to them. out = list(out) # Get + merge parts. parts = [spec.instance_from_str(arg_list) for arg_list in x] for part in parts: if isinstance(out, dict): out.update(part) else: out.append(part) # Return output with correct type. if container_type in (dict, Sequence, collections.abc.Sequence): return out else: return container_type(out) lowered.instance_from_str = append_instantiator lowered.str_from_instance = spec.str_from_instance lowered.choices = spec.choices lowered.nargs = spec.nargs if not isinstance(spec.nargs, tuple) else "*" lowered.metavar = spec.metavar lowered.action = spec._action lowered.required = False return else: lowered.instance_from_str = spec.instance_from_str lowered.str_from_instance = spec.str_from_instance lowered.choices = spec.choices lowered.nargs = spec.nargs if not isinstance(spec.nargs, tuple) else "*" lowered.metavar = spec.metavar lowered.action = spec._action return def _rule_counters( arg: ArgumentDefinition, lowered: LoweredArgumentDefinition, ) -> None: """Handle counters, like -vvv for level-3 verbosity.""" if ( _markers.UseCounterAction in arg.field.markers and arg.field.type_stripped is int and not arg.is_positional() ): lowered.metavar = None lowered.nargs = None lowered.action = "count" lowered.default = ( arg.field.default if not _singleton.is_missing(arg.field.default) else 0 ) lowered.required = False lowered.instance_from_str = lambda x: ( x ) # argparse will directly give us an int! return def _rule_generate_helptext( arg: ArgumentDefinition, lowered: LoweredArgumentDefinition, ) -> None: """Store context for lazy helptext generation.""" # Defer helptext generation by wrapping in a callable. # The percent symbol needs some extra handling in argparse. # https://stackoverflow.com/questions/21168120/python-argparse-errors-with-in-help-string lowered.help = lambda: ( generate_argument_helptext(arg, lowered).as_str_no_ansi().replace("%", "%%") ) def _rule_set_name_or_flag_and_dest( arg: ArgumentDefinition, lowered: LoweredArgumentDefinition, ) -> None: extern_name = arg.field.extern_name if lowered.action == "store_false": extern_name = "no_" + extern_name if ( arg.field.argconf.prefix_name is False or _markers.OmitArgPrefixes in arg.field.markers ): # Strip prefixes when the argument is suppressed. # Still need to call make_field_name() because it converts underscores # to hyphens, etc. name_or_flag = _strings.make_field_name([extern_name]) elif ( _markers.OmitSubcommandPrefixes in arg.field.markers and arg.subcommand_prefix != "" ): # Strip subcommand prefixes, but keep following # prefixes.`extern_prefix` can start with the prefix corresponding to # the parent subcommand, but end with other prefixes corresponding to # nested structures within the subcommand. name_or_flag = _strings.make_field_name([arg.extern_prefix, extern_name]) strip_prefix = arg.subcommand_prefix + "." assert name_or_flag.startswith(strip_prefix), name_or_flag name_or_flag = name_or_flag[len(strip_prefix) :] else: # Standard prefixed name. name_or_flag = _strings.make_field_name([arg.extern_prefix, extern_name]) # Prefix keyword arguments with --. if not arg.is_positional(): name_or_flag = "--" + name_or_flag lowered.name_or_flags = (name_or_flag,) lowered.dest = _strings.make_field_name([arg.intern_prefix, arg.field.intern_name]) def _rule_positional_special_handling( arg: ArgumentDefinition, lowered: LoweredArgumentDefinition, ) -> None: if not arg.is_positional(): return None metavar = lowered.metavar # Positional arguments with nargs="*" accept zero arguments, so they # should never be marked as required. if lowered.required and lowered.nargs == "*": lowered.required = False if lowered.required: nargs = lowered.nargs else: if metavar is not None: metavar = "[" + metavar + "]" if lowered.nargs == 1: # Optional positional arguments. This needs to be special-cased in # _calling.py. nargs = "?" else: # If lowered.nargs is either + or an int. nargs = "*" lowered.name_or_flags = ( _strings.make_field_name([arg.intern_prefix, arg.field.intern_name]), ) lowered.dest = None lowered.metavar = metavar lowered.nargs = nargs return def _rule_apply_argconf( arg: ArgumentDefinition, lowered: LoweredArgumentDefinition, ) -> None: if arg.field.argconf.metavar is not None: lowered.metavar = arg.field.argconf.metavar if arg.field.argconf.aliases is not None: lowered.name_or_flags = arg.field.argconf.aliases + lowered.name_or_flags def generate_argument_helptext( arg: ArgumentDefinition, lowered: LoweredArgumentDefinition, compact: bool = False ) -> fmt._Text: help_parts: list[str | fmt._Text] = [] primary_help = arg.field.helptext # Evaluate lazy helptext if it's a callable. if callable(primary_help): primary_help = primary_help() if primary_help is None and _markers.Positional in arg.field.markers: primary_help = _strings.make_field_name( [arg.extern_prefix, arg.field.intern_name] ) # In compact mode, skip the primary help text. if primary_help is not None and not compact: help_parts.append(fmt.text["dim"](primary_help)) if not lowered.required: # Get the default value. # Note: lowered.default is the stringified version! if ( arg.is_positional() and lowered.nargs == "*" and _singleton.is_missing(arg.field.default) ): # Positional arguments with nargs="*" and no explicit default # effectively default to an empty sequence. default = () elif ( lowered.is_fixed() or lowered.action in ( "append", "boolean_optional_action", "store_true", "store_false", ) or _singleton.is_missing(arg.field.default) ): # Cases where we want to use the field default directly. default = arg.field.default elif arg.field.default is _singleton.EXCLUDE_FROM_CALL: default = None else: # Standard cases where we convert to a string representation. default = ( lowered.str_from_instance(arg.field.default) if lowered.str_from_instance is not None else None ) # Get the default value label. if arg.field.argconf.constructor_factory is not None: default_label = ( str(default) if arg.field.type_stripped is not json.loads else json.dumps(arg.field.default) ) elif type(default) in (tuple, list, set): # For tuple types, we might have default as (0, 1, 2, 3). # For list types, we might have default as [0, 1, 2, 3]. # For set types, we might have default as {0, 1, 2, 3}. # # In all cases, we want to display (default: 0 1 2 3), for consistency with # the format that argparse expects when we set nargs. assert default is not None default_label = " ".join(map(shlex.quote, map(str, default))) else: default_label = str(default) # Suffix helptext with some behavior hint, such as the default value of the argument. help_behavior_hint = arg.field.argconf.help_behavior_hint if help_behavior_hint is not None: behavior_hint = ( help_behavior_hint(default_label) if callable(help_behavior_hint) else help_behavior_hint ) elif lowered.instance_from_str is None: # Intentionally not quoted via shlex, since this can't actually be passed # in via the commandline. behavior_hint = f"(fixed to: {str(default)})" elif lowered.action == "count": # Repeatable argument. behavior_hint = "(repeatable)" elif lowered.action == "append" and ( _singleton.is_missing(default) or len(cast(tuple, default)) == 0 ): behavior_hint = "(repeatable)" elif lowered.action == "append" and len(cast(tuple, default)) > 0: assert default is not None # Just for type checker. behavior_hint = f"(repeatable, appends to: {default_label})" elif arg.field.default is _singleton.EXCLUDE_FROM_CALL: # ^important to use arg.field.default and not the stringified default variable. behavior_hint = "(unset by default)" elif _markers._OPTIONAL_GROUP in arg.field.markers and _singleton.is_missing( default ): # Argument in an optional group, but with no default. This is typically used # when general (non-argument, non-dataclass) object arguments are given a # default, or when we use `tyro.conf.arg(constructor=...)`. # # There are some usage details that aren't communicated right now in the # helptext. For example: all arguments within an optional group without a # default should be passed in or none at all. behavior_hint = "(optional)" elif _markers._OPTIONAL_GROUP in arg.field.markers: # Argument in an optional group, but which also has a default. behavior_hint = f"(default if used: {default_label})" else: behavior_hint = f"(default: {default_label})" help_parts.append( fmt.text[ _settings.ACCENT_COLOR if _settings.ACCENT_COLOR != "white" else "cyan" ](behavior_hint) ) else: help_parts.append(fmt.text["bright_red"]("(required)")) return fmt.text(*help_parts, delimiter=" ")