test_choice.py

from dataclasses import dataclass, field
from enum import Enum
from typing import Union

import pytest

from simple_parsing import choice

from .testutils import TestSetup, raises


@dataclass
class A(TestSetup):
    color: str = choice("red", "green", "blue", default="red")
    colors: list[str] = choice("red", "green", "blue", default_factory=["red"].copy)


def test_choice_default():
    a = A.setup("")
    assert a.color == "red"
    assert a.colors == ["red"]


def test_value_not_in_choices_throws_error():
    with raises(SystemExit):
        A.setup("--color orange")
    with raises(SystemExit):
        A.setup("--colors red orange")


def test_passed_value_works_fine():
    a = A.setup("--color red")
    assert a.color == "red"

    a = A.setup("--color green")
    assert a.color == "green"

    a = A.setup("--color blue")
    assert a.color == "blue"

    a = A.setup("--colors red blue")
    assert a.colors == ["red", "blue"]

    a = A.setup("--colors blue red green red")
    assert a.colors == ["blue", "red", "green", "red"]


@dataclass
class Base:
    value: str = "hello base"


@dataclass
class AA(Base):
    value: str = "hello a"


@dataclass
class BB(Base):
    value: str = "hello b"


def test_choice_with_dict():
    @dataclass
    class C(TestSetup):
        option: Union[AA, BB, float] = choice(
            {"a": AA("aaa"), "b": BB("bbb"), "bob": AA("bobobo"), "f": 1.23},
            default="a",
        )
        options: list[Union[AA, BB, float]] = choice(
            {"a": AA("aaa"), "b": BB("bbb"), "bob": AA("bobobo"), "f": 1.23},
            default_factory=["a"].copy,
        )

    c = C.setup("")
    assert c.option == AA("aaa")
    assert c.options == [AA("aaa")]

    c = C.setup("--option a --options a a")
    assert c.option == AA("aaa")
    assert c.options == [AA("aaa"), AA("aaa")]

    c = C.setup("--option bob --options bob a")
    assert c.option == AA("bobobo")
    assert c.options == [AA("bobobo"), AA("aaa")]

    c = C.setup("--option f --options a f a b")
    assert c.option == 1.23
    assert c.options == [AA("aaa"), 1.23, AA("aaa"), BB("bbb")]


def test_choice_with_default_instance():
    @dataclass
    class D(TestSetup):
        option: Union[AA, BB, float] = choice(
            {
                "a": [AA("aa1"), AA("aa2")],
                "b": 1.23,
                "bob": BB("bobobo"),
            },
            default="a",
        )

    @dataclass
    class Parent(TestSetup):
        d: D = field(default_factory=lambda: D(option=AA("parent")))

    p = Parent.setup("")
    assert p.d.option == AA("parent")


class Color(Enum):
    blue: str = "BLUE"
    red: str = "RED"
    green: str = "GREEN"
    orange: str = "ORANGE"


def test_passing_enum_to_choice():
    @dataclass
    class Something(TestSetup):
        favorite_color: Color = choice(Color, default=Color.green)
        colors: list[Color] = choice(Color, default_factory=[Color.green].copy)

    s = Something.setup("")
    assert s.favorite_color == Color.green
    assert s.colors == [Color.green]

    s = Something.setup("--colors blue red")
    assert s.colors == [Color.blue, Color.red]


# TODO: Add a test to check how the enum arguments are shown on the help string.


def test_passing_enum_to_choice_no_default_makes_required_arg():
    @dataclass
    class Something(TestSetup):
        favorite_color: Color = choice(Color)

    with raises(SystemExit):
        s = Something.setup("")

    s = Something.setup("--favorite_color blue")
    assert s.favorite_color == Color.blue


def test_passing_enum_to_choice_with_key_as_default():
    with pytest.warns(UserWarning):

        @dataclass
        class Something(TestSetup):
            favorite_color: Color = choice(Color, default="blue")


def test_passing_enum_to_choice_is_same_as_enum_attr():
    @dataclass
    class Something1(TestSetup):
        favorite_color: Color = Color.orange

    @dataclass
    class Something2(TestSetup):
        favorite_color: Color = choice(Color, default=Color.orange)

    s1 = Something1.setup("--favorite_color green")
    s2 = Something2.setup("--favorite_color green")
    assert s1.favorite_color == s2.favorite_color

    s = Something1.setup("--favorite_color blue")
    assert s.favorite_color == Color.blue
    s = Something2.setup("--favorite_color blue")
    assert s.favorite_color == Color.blue