def proto_to_choice_space()

in compiler_gym/service/proto/py_converters.py [0:0]

• 50 lines of code
• 9 McCabe index (conditional complexity)

def proto_to_choice_space(choice: ChoiceSpace) -> PyChoiceSpace:
    """Convert a ChoiceSpace message to a gym.Space and choice callback.

    :param proto: A ChoiceSpace message.

    :returns: A PyChoiceSpace tuple, comprising a gym.Space space, and a
        callback that wraps an input in a Choice message.
    """
    choice_type = choice.WhichOneof("space")
    if choice_type == "int64_range":
        # The Discrete class defines a discrete space as integers in the range
        # [0,n]. For spaces that aren't zero-based there is a Scalar class.
        # Prefer Discrete if possible since it is part of the core gym library,
        # else fallback to Scalar.
        if choice.int64_range.min.value:
            return PyChoiceSpace(
                space=Scalar(
                    min=choice.int64_range.min.value,
                    max=choice.int64_range.max.value,
                    dtype=int,
                    name=choice.name,
                ),
                make_choice=lambda a: Choice(int64_value=a),
            )
        else:
            return PyChoiceSpace(
                space=Discrete(n=choice.int64_range.max.value, name=choice.name),
                make_choice=lambda a: Choice(int64_value=a),
            )
    elif choice_type == "double_range":
        return PyChoiceSpace(
            space=Scalar(
                min=choice.double_range.min.value,
                max=choice.double_range.max.value,
                dtype=float,
                name=choice.name,
            ),
            make_choice=lambda a: Choice(double_value=a),
        )
    elif (
        choice_type == "named_discrete_space"
        and choice.named_discrete_space.is_commandline
    ):
        return PyChoiceSpace(
            space=Commandline(
                items=[
                    CommandlineFlag(name=x, flag=x, description="")
                    for x in choice.named_discrete_space.value
                ],
                name=choice.name,
            ),
            make_choice=lambda a: Choice(named_discrete_value_index=a),
        )
    elif choice_type == "named_discrete_space":
        return PyChoiceSpace(
            space=NamedDiscrete(
                items=choice.named_discrete_space.value, name=choice.name
            ),
            make_choice=lambda a: Choice(named_discrete_value_index=a),
        )
    raise ValueError(f"Invalid ChoiceSpace: {choice}")