Support Enum aliases (#162)

* populate enum aliases * add test * autogenerated test * Update test_dcargs.py Signed-off-by: Brent Yi <yibrenth@gmail.com> * Update test_dcargs_generated.py Signed-off-by: Brent Yi <yibrenth@gmail.com> --------- Signed-off-by: Brent Yi <yibrenth@gmail.com> Co-authored-by: Brent Yi <yibrenth@gmail.com>
brentyi · Sep 14, 2024 · bf3f61a · bf3f61a
1 parent 072dd5f
commit bf3f61a
Show file tree

Hide file tree

Showing 3 changed files with 30 additions and 1 deletion.
diff --git a/src/tyro/_instantiators.py b/src/tyro/_instantiators.py
@@ -277,7 +277,8 @@ def instantiate(args: List[str]) -> Any:
     if typ is bool:
         auto_choices = ("True", "False")
     elif inspect.isclass(typ) and issubclass(typ, enum.Enum):
-        auto_choices = tuple(x.name for x in typ)
+        # Using `.__members__` dict to handle aliases correctly
+        auto_choices = tuple(typ.__members__.keys())
 
     def instantiator_base_case(strings: List[str]) -> Any:
         """Given a type and and a string from the command-line, reconstruct an object. Not

diff --git a/tests/test_dcargs.py b/tests/test_dcargs.py
@@ -297,6 +297,20 @@ class EnumClassB:
     assert tyro.cli(EnumClassB, args=[]) == EnumClassB()
 
 
+def test_enum_alias() -> None:
+    class Color(enum.Enum):
+        RED = 1
+        ROUGE = 1
+
+    @dataclasses.dataclass
+    class A:
+        color: Color
+
+    assert tyro.cli(A, args=["--color", "RED"]) == A(color=Color.RED)
+    assert tyro.cli(A, args=["--color", "ROUGE"]) == A(color=Color.ROUGE)
+    assert tyro.cli(A, args=["--color", "ROUGE"]) == A(color=Color.RED)
+
+
 def test_literal() -> None:
     @dataclasses.dataclass
     class A:

diff --git a/tests/test_py311_generated/test_dcargs_generated.py b/tests/test_py311_generated/test_dcargs_generated.py
@@ -299,6 +299,20 @@ class EnumClassB:
     assert tyro.cli(EnumClassB, args=[]) == EnumClassB()
 
 
+def test_enum_alias() -> None:
+    class Color(enum.Enum):
+        RED = 1
+        ROUGE = 1
+
+    @dataclasses.dataclass
+    class A:
+        color: Color
+
+    assert tyro.cli(A, args=["--color", "RED"]) == A(color=Color.RED)
+    assert tyro.cli(A, args=["--color", "ROUGE"]) == A(color=Color.ROUGE)
+    assert tyro.cli(A, args=["--color", "ROUGE"]) == A(color=Color.RED)
+
+
 def test_literal() -> None:
     @dataclasses.dataclass
     class A: