diff --git a/simple_parsing/wrappers/field_parsing.py b/simple_parsing/wrappers/field_parsing.py index 5208a6ec..493c4b3c 100644 --- a/simple_parsing/wrappers/field_parsing.py +++ b/simple_parsing/wrappers/field_parsing.py @@ -16,6 +16,7 @@ is_forward_ref, is_homogeneous_tuple_type, is_list, + is_literal, is_tuple, is_typevar, is_union, @@ -127,6 +128,10 @@ def get_parsing_fn(t: type[T]) -> Callable[[Any], T]: args = get_type_arguments(t) return parse_union(*args) + elif is_literal(t): + logger.debug(f"Parsing a Literal field of type {t}") + return parse_literal(t) + elif is_enum(t): logger.debug(f"Parsing an Enum field of type {t}") return parse_enum(t) @@ -205,6 +210,29 @@ def _parse_optional(val: Optional[Any]) -> Optional[T]: return _parse_optional +def parse_literal(literal_type: type[T]) -> Callable[[str], T]: + """Returns a parsing function for a Literal type. + + The function maps the string representation of each literal value back to the actual value + (e.g. "1" -> 1, "BLUE" -> Color.BLUE). + """ + literal_values = get_type_arguments(literal_type) + # Build a mapping from the string representation to the actual value. + choice_dict: dict[str, Any] = { + (v.name if isinstance(v, enum.Enum) else str(v)): v for v in literal_values + } + + def _parse_literal(val: str) -> T: + if val in choice_dict: + return choice_dict[val] + raise ValueError( + f"Invalid value {val!r} for {literal_type}. Expected one of: {list(choice_dict)}" + ) + + _parse_literal.__name__ = str(literal_type) + return _parse_literal + + def parse_tuple(tuple_item_types: tuple[type[T], ...]) -> Callable[[list[T]], tuple[T, ...]]: """Makes a parsing function for creating tuples from the command-line args. diff --git a/test/test_literal.py b/test/test_literal.py index b251056b..8fa6b186 100644 --- a/test/test_literal.py +++ b/test/test_literal.py @@ -1,7 +1,7 @@ import enum import sys from dataclasses import dataclass -from typing import Any, NamedTuple, Optional +from typing import Any, NamedTuple, Optional, Union import pytest from typing_extensions import Literal @@ -118,3 +118,21 @@ def test_reproduce_issue_259_parsing_literal_py39(): "argument --param: invalid typing.Literal['bar', 'biz'] value: 'biz'" ): assert SomeFoo.setup("").param == "biz" + + +@dataclass +class Foo: + bar: Union[Literal["a"], int] = "a" + + +def test_issue_322(): + """Test for https://github.com/lebrice/SimpleParsing/issues/322.""" + from simple_parsing import parse + + assert parse(Foo, args="") == Foo() + assert parse(Foo, args="--bar=a") == Foo(bar="a") + # 'b' is neither in Literal["a"] nor a valid int, so it should be rejected. + with exits_and_writes_to_stderr("invalid"): + assert parse(Foo, args="--bar=b") + + assert parse(Foo, args="--bar=123") == Foo(bar=123)