Allow choices restriction

[Borda]

Hello, and thank you for this great CLI!
Recently I get to a situation when I would like to restrict the options for a given argument similar to build-in argparse does with its option choices (see docs: https://docs.python.org/3/library/argparse.html#choices). Then I was checking Fire docs but could not find anything similar to it...
Checking alternative CLI packages I found a way that is quite simple but still elegant and would well fit the Fire style. It is leveraging python Enum class:

from enum import Enum
import fire

class Direction(str, Enum):
    up = "up"
    down = "down"
    left = "left"
    right = "right"


def main(move: Direction = Direction.left):
    print(f"Moving in given direction: {move.value}")


if __name__ == "__main__":
    fire.Fire(main)

For clarification, the example above is borrowed and adjusted from Typer/enum

[chris-clem]

That would be great! I use choices quite a lot.

[narothsolo]

Please help Star

[dbieber]

Great idea. We don't currently use type annotations in fire to impose restrictions (but we could in a future version, though no one is actively working toward it atm).

Side note: One alternative that works today is to use a decorator, roughly like this:

def restrict_choices(choices):
  def decorator(f):
    def new_f(x):
      if x not in choices:
        raise FireError("Invalid choice")
      return f(x)
    return new_f
  return decorator

@restrict_choices(['left', 'right'])
def main(move):
    print(f"Moving in given direction: {move}")

See also SetParseFns in https://github.com/google/python-fire/blob/master/fire/decorators.py

[keyboardAnt]

You might also find the HfArgumentParser relevant: https://github.com/huggingface/transformers/blob/514de24abfd4416aeba6a6455ad5920f57f3567d/src/transformers/hf_argparser.py#L109

[Borda]

You might also find the HfArgumentParser relevant: https://github.com/huggingface/transformers/blob/514de24abfd4416aeba6a6455ad5920f57f3567d/src/transformers/hf_argparser.py#L109

Not really if you have to install full HF package for it...

[keyboardAnt]

You might also find the HfArgumentParser relevant: https://github.com/huggingface/transformers/blob/514de24abfd4416aeba6a6455ad5920f57f3567d/src/transformers/hf_argparser.py#L109

Not really if you have to install full HF package for it...

The alternative below doesn't need the HF package. It is simple and readable but creates the Config object twice.

from pydantic import BaseModel

class Config(BaseModel):
    ...

def main(**kwargs):
    config = Config().model_copy(update=kwargs)

if __name__ == "__main__":
    fire.Fire(main)

[hesic73]

Thanks! Here is a more generalized version with the help of GPT. The problem is that if sig.bind raises an TypeError Exception, fire won't work.

from typing import Union, List, Any
import inspect
from fire.core import FireError


def restrict_choices(arg_name_or_position: Union[int, str], choices: List[Any]):
    def decorator(f):
        sig = inspect.signature(f)  # Get the function signature

        def new_f(*args, **kwargs):
            # Map arguments by position and name
            bound_args = sig.bind(*args, **kwargs)
            bound_args.apply_defaults()  # Handle any default arguments

            # Determine if we're restricting by name or position
            if isinstance(arg_name_or_position, str):
                # Restrict by argument name
                if arg_name_or_position in bound_args.arguments:
                    restricted_arg = bound_args.arguments[arg_name_or_position]
                    arg_identifier = f"argument '{arg_name_or_position}'"
                else:
                    raise FireError(
                        f"Argument '{arg_name_or_position}' not found")
            elif isinstance(arg_name_or_position, int):
                # Restrict by argument position
                if arg_name_or_position < len(bound_args.args):
                    restricted_arg = bound_args.args[arg_name_or_position]
                    arg_identifier = f"position {arg_name_or_position}"
                else:
                    raise FireError(
                        f"Argument position {arg_name_or_position} is out of range")
            else:
                raise FireError(
                    "Invalid argument specifier, must be a name (str) or position (int)")

            # Check if the restricted argument is in the allowed choices
            if restricted_arg not in choices:
                raise FireError(
                    f"Invalid choice '{restricted_arg}' for {arg_identifier}. "
                    f"Valid choices are: {choices}")

            # Call the original function if the check passes
            return f(*args, **kwargs)

        return new_f
    return decorator

Examples:

@restrict_choices('direction', ['left', 'right'])
def move(direction, speed):
    print(f"Moving {direction} at speed {speed}")


@restrict_choices(1, ['left', 'right'])
def move(speed, direction):
    print(f"Moving {direction} at speed {speed}")

[rgon]

To solve this temporarily, without manually configuring each decorator, you can use this:

@FireTyped
async def yourfn(...)

For this to not require any manual config (as is fire philosophy), I had to implement full argument type handling. So it'll parse the type of all arguments passed, which has the added benefit of restricting enums. So yourfn(number:int) will get type(number) === int as well, and yourfn(id:str) when called with --id 45 will now correctly get the str type (#327 #453).

It also handles multiple enum types YourEnum(str, Enum) and YourNumberEnum(int, Enum) by offering to select their names, instead of their values.

It also allows defining a from_string method (like datetime has) to manually parse types. If not, it calls the constructor.

(mainlining this would close #300 #327 #453 #260). Similar to PRs #350 (but this supports arbitrary/all types) and the closed #320 . Related: #508

import inspect
from enum import Enum
from termcolor import cprint
from fire.decorators import SetParseFns
# from fire.value_types import VALUE_TYPES
from typing import Any, Type, Callable, Annotated

def FireTyped(func:Callable):
    """
    Ensures that all arguments of the called function are validated against their type annotations,
    and are converted to the correct type.

    The type's constructor will be passed the value from the command line (parsed by Fire, so: str|int|float|bool|tuple|list|dict).

    Decorator that uses Fire's SetParseFn to ensure all annotated arguments are parsed to their type,
    including Enums, and prints helpful errors on failure.
    """

    def parse_arg(value:Annotated[Any, "type of fire VALUE_TYPES"], annotation:Type):
        if value is None:
            return value
        try:
            if hasattr(annotation, 'from_string'):
                # If the annotation has a from_string method, call it with the value
                # Example: datetime.from_string(value)
                return annotation.from_string(value)

            elif isinstance(annotation, type) and issubclass(annotation, Enum):
                # If the annotation is an Enum, convert the value to the Enum member
                # Try to get from key first
                try:
                    return annotation[value]
                except KeyError:
                    raise ValueError(f"Invalid Enum key name: {value}")

            elif isinstance(annotation, type):
                return annotation(value)
            else:
                return value

        except ValueError as e:
            import sys
            print(f"Could not convert value '{value}' to {annotation.__name__}:")
            cprint(e, "red")

            if isinstance(annotation, type) and issubclass(annotation, Enum):
                cprint(f"Valid options:", "yellow")
                for option in annotation:
                    cprint(f"  {option.name}: {option.value}", "yellow")

            # Exit gracefully
            sys.exit(64) # from sysexits.h - EX_USAGE 
        except Exception as e:
            cprint(f"Unexpected error occurred while converting value '{value}' to {annotation.__name__}: {e}", "red")
            raise e

    sig = inspect.signature(func)
    named_type_map = {}
    for name, param in sig.parameters.items():
        annotation = param.annotation
        if annotation != inspect.Parameter.empty:
            # Capture the value of the lambda arguments with = to ensure they're passed correctly to parse_arg
            named_type_map[name] = lambda v, annotation=annotation: parse_arg(v, annotation)

    # Use the built-in SetParseFns decorator
    decorator = SetParseFns(**named_type_map)
    return decorator(func)