def proto_to_action_space()

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

• 36 lines of code
• 10 McCabe index (conditional complexity)

def proto_to_action_space(proto: ActionSpace) -> PyActionSpace:
    """Convert a ActionSpace message to a gym.Space and action callback.

    :param proto: An ActionSpace message.

    :returns: A PyActionSpace tuple, comprising a gym.Space space, and a
        callback that formats its input as an Action message.
    """
    if proto.named_choices:
        # Convert a list of named choices to a dictionary space.
        choices = [proto_to_choice_space(choice) for choice in proto.choice]

        def compose_choice_dict(action: DictType[str, Any]):
            return Action(
                choice=[
                    choice.make_choice(action[choice.space.name]) for choice in choices
                ]
            )

        return PyActionSpace(
            space=Dict(
                spaces={
                    choice.name: space.space
                    for choice, space in zip(proto.choice, choices)
                },
                name=proto.name,
            ),
            make_action=compose_choice_dict,
        )
    elif len(proto.choice) > 1:
        # Convert an unnamed list of choices to a tuple space.
        choices = [proto_to_choice_space(choice) for choice in proto.choice]

        def compose_choice_list(action: List[Any]) -> Action:
            return Action(
                choice=[choice.make_choice(a) for a, choice in zip(action, choices)]
            )

        return PyActionSpace(
            space=Tuple(spaces=[choice.space for choice in choices], name=proto.name),
            make_action=compose_choice_list,
        )
    elif proto.choice:
        # Convert a single choice into an action space.
        space, make_choice = proto_to_choice_space(proto.choice[0])
        # When there is only a single choice, use the name of the parent space.
        space.name = proto.name
        return PyActionSpace(
            space=space, make_action=lambda a: Action(choice=[make_choice(a)])
        )
    raise ValueError("No choices set for ActionSpace")