Support typing.Literal in config flags.

PiperOrigin-RevId: 861237487

Author: The ml_collections Authors

ml_collections/config_flags/config_flags.py

@@ -27,6 +27,7 @@
 import sys
 import traceback
 import types
+import typing
 from typing import Any, Callable, Dict, Generic, List, MutableMapping, Optional, Sequence, Tuple, Type, TypeVar
 
 from absl import flags
@@ -859,7 +860,10 @@ def _parse(self, argument):
         parser = _FIELD_TYPE_TO_PARSER[config_dict.ConfigDict]
       elif field_type_origin and field_type_origin in _FIELD_TYPE_TO_PARSER:
         parser = _FIELD_TYPE_TO_PARSER[field_type_origin]
-      elif issubclass(field_type, enum.Enum):
+      elif field_type_origin is typing.Literal:
+        # Literal types like Literal["a", "b"] should be treated as string-like.
+        parser = _LiteralParser()
+      elif isinstance(field_type, type) and issubclass(field_type, enum.Enum):
         parser = flags.EnumClassParser(field_type, case_sensitive=False)
       elif dataclasses.is_dataclass(field_type):
         # For dataclasses-valued fields allow default instance creation.

ml_collections/config_flags/tests/dataclass_overriding_test.py

@@ -18,7 +18,7 @@
 import dataclasses
 import functools
 import sys
-from typing import Mapping, Optional, Sequence, Tuple, Union
+from typing import Literal, Mapping, Optional, Sequence, Tuple, Union
 import unittest
 
 from absl import flags
@@ -179,6 +179,17 @@ class PipeConfig:
     result = _test_flags(PipeConfig(), '.foo=32')
     self.assertEqual(result.foo, 32)
 
+  def test_literal_type_field(self):
+    @dataclasses.dataclass
+    class ConfigWithLiteral:
+      optimizer: Literal['adam', 'sgd', 'rmsprop'] = 'adam'
+
+    result = _test_flags(ConfigWithLiteral())
+    self.assertEqual(result.optimizer, 'adam')
+
+    result = _test_flags(ConfigWithLiteral(), '.optimizer=sgd')
+    self.assertEqual(result.optimizer, 'sgd')
+
   def test_custom_flag_parsing_override_work(self):
     # Overrides still work.
     result = _test_flags(_CONFIG, '.custom.i=10')