Always use Literal evals when parsing kauldron flags

PiperOrigin-RevId: 872335743

Author: Conchylicultor

ml_collections/config_flags/__init__.py

@@ -19,6 +19,7 @@
 from .config_flags import DEFINE_config_file
 from .config_flags import get_config_filename
 from .config_flags import get_override_values
+from .config_flags import OverrideMode
 from .config_flags import register_flag_parser
 from .config_flags import register_flag_parser_for_type
 
@@ -28,6 +29,7 @@
     "DEFINE_config_file",
     "get_config_filename",
     "get_override_values",
+    "OverrideMode",
     "register_flag_parser",
     "register_flag_parser_for_type",
 )

ml_collections/config_flags/config_flags.py

@@ -15,6 +15,7 @@
 """Configuration commmand line parser."""
 
 import ast
+import contextlib
 import copy
 import dataclasses
 import enum
@@ -48,6 +49,41 @@
 flags.disclaim_key_flags()
 
 
+class OverrideMode(enum.Enum):
+  """Behavior to overwrite config parameters types.
+
+  Attributes:
+    ALWAYS: All `--cfg.xxx` flags are parsed using `_LiteralParser`.
+    NEW_ONLY: `--cfg.xxx` is only parsed using `_LiteralParser`, if `cfg.xxx`
+      does not exist yet.
+    NEVER: (default) All `--cfg.xxx` flags are parsed using type-specific
+      parsers (i.e. it is not possible to have flag supporting both int and
+      str).
+  """
+
+  ALWAYS = enum.auto()
+  NEW_ONLY = enum.auto()
+  NEVER = enum.auto()
+
+
+def _normalize_accept_new_attributes(
+    accept_new_attributes: bool | None = None,
+    override_mode: OverrideMode | None = None,
+) -> OverrideMode:
+  """Normalizes override_mode."""
+  if accept_new_attributes is not None and override_mode is not None:
+    raise ValueError(
+        'Cannot specify both `accept_new_attributes` and `override_mode`.'
+    )
+  if override_mode is not None:
+    return override_mode
+  if accept_new_attributes is None:
+    return OverrideMode.NEVER
+  if accept_new_attributes:
+    return OverrideMode.NEW_ONLY
+  return OverrideMode.NEVER
+
+
 def _load_source(module_name: str, module_path: str) -> types.ModuleType:
   """Loads a Python module from its source file.
 
@@ -145,9 +181,11 @@ def DEFINE_config_file(  # pylint: disable=g-bad-name
     help_string: str = 'path to config file.',
     flag_values: flags.FlagValues = FLAGS,
     lock_config: bool = True,
-    accept_new_attributes: bool = False,
+    accept_new_attributes: bool | None = None,
+    override_mode: OverrideMode | None = None,
     sys_argv: Optional[List[str]] = None,
-    **kwargs) -> flags.FlagHolder:
+    **kwargs,
+) -> flags.FlagHolder:
   r"""Defines flag for `ConfigDict` files compatible with absl flags.
 
   The flag's value should be a path to a valid python file which contains a
@@ -247,9 +285,10 @@ def get_config(config_string):
       absl.flags.FLAGS)
     lock_config: If set to True, loaded config will be locked through calling
       .lock() method on its instance (if it exists). (default: True)
-    accept_new_attributes: If `True`, accept to pass arbitrary attributes that
-      are not originally defined in the `get_config()` dict.
-      `accept_new_attributes` requires `lock_config=False`.
+    accept_new_attributes: Deprecated, use `override_mode` instead. `False`:
+      only existing attributes can be overridden. `True`: new attributes are
+      accepted. Requires `lock_config=False`.
+    override_mode: Controls how override flags are parsed.
     sys_argv: If set, interprets this as the full list of args used in parsing.
       This is used to identify which overrides to define as flags. If not
       specified, uses the system sys.argv to figure it out.
@@ -258,8 +297,12 @@ def get_config(config_string):
   Returns:
     a handle to defined flag.
   """
-  if accept_new_attributes and lock_config:
-    raise ValueError('`accept_new_attributes=True` requires lock_config=False')
+
+  override_mode = _normalize_accept_new_attributes(
+      accept_new_attributes, override_mode
+  )
+  if override_mode != OverrideMode.NEVER and lock_config:
+    raise ValueError('`override_mode` requires `lock_config=False`')
   parser = ConfigFileFlagParser(name=name, lock_config=lock_config)
   serializer = flags.ArgumentSerializer()
   flag = _ConfigFlag(
@@ -269,9 +312,10 @@ def get_config(config_string):
       default=default,
       help_string=help_string,
       flag_values=flag_values,
-      accept_new_attributes=accept_new_attributes,
+      override_mode=override_mode,
       sys_argv=sys_argv,
-      **kwargs)
+      **kwargs,
+  )
 
   return flags.DEFINE_flag(flag, flag_values)
 
@@ -282,9 +326,11 @@ def DEFINE_config_dict(  # pylint: disable=g-bad-name
     help_string: str = 'ConfigDict instance.',
     flag_values: flags.FlagValues = FLAGS,
     lock_config: bool = True,
-    accept_new_attributes: bool = False,
+    accept_new_attributes: bool | None = None,
+    override_mode: OverrideMode | None = None,
     sys_argv: Optional[List[str]] = None,
-    **kwargs) -> flags.FlagHolder:
+    **kwargs,
+) -> flags.FlagHolder:
   """Defines flag for inline `ConfigDict's` compatible with absl flags.
 
   Similar to `DEFINE_config_file` except the flag's value should be a
@@ -328,15 +374,16 @@ def DEFINE_config_dict(  # pylint: disable=g-bad-name
   Args:
     name: Flag name.
     config: `ConfigDict` object.
-    help_string: Help string to display when --helpfull is called.
-        (default: "ConfigDict instance.")
-    flag_values: FlagValues instance used for parsing.
-        (default: absl.flags.FLAGS)
+    help_string: Help string to display when --helpfull is called. (default:
+      "ConfigDict instance.")
+    flag_values: FlagValues instance used for parsing. (default:
+      absl.flags.FLAGS)
     lock_config: If set to True, loaded config will be locked through calling
-        .lock() method on its instance (if it exists). (default: True)
-    accept_new_attributes: If `True`, accept to pass arbitrary attributes that
-      are not originally defined in the `config` argument.
-      `accept_new_attributes` requires `lock_config=False`.
+      .lock() method on its instance (if it exists). (default: True)
+    accept_new_attributes: Deprecated, use `override_mode` instead. `False`:
+      only existing attributes can be overridden. `True`: new attributes are
+      accepted. Requires `lock_config=False`.
+    override_mode: Controls how override flags are parsed.
     sys_argv: If set, interprets this as the full list of args used in parsing.
       This is used to identify which overrides to define as flags. If not
       specified, uses the system sys.argv to figure it out.
@@ -347,8 +394,11 @@ def DEFINE_config_dict(  # pylint: disable=g-bad-name
   """
   if not isinstance(config, config_dict.ConfigDict):
     raise TypeError('config should be a ConfigDict')
-  if accept_new_attributes and lock_config:
-    raise ValueError('`accept_new_attributes=True` requires lock_config=False')
+  override_mode = _normalize_accept_new_attributes(
+      accept_new_attributes, override_mode
+  )
+  if override_mode != OverrideMode.NEVER and lock_config:
+    raise ValueError('`override_mode` requires `lock_config=False`')
   parser = _InlineConfigParser(name=name, lock_config=lock_config)
   flag = _ConfigFlag(
       parser=parser,
@@ -357,9 +407,10 @@ def DEFINE_config_dict(  # pylint: disable=g-bad-name
       default=config,
       help_string=help_string,
       flag_values=flag_values,
-      accept_new_attributes=accept_new_attributes,
+      override_mode=override_mode,
       sys_argv=sys_argv,
-      **kwargs)
+      **kwargs,
+  )
 
   # Get the module name for the frame at depth 1 in the call stack.
   module_name = sys._getframe(1).f_globals.get('__name__', None)  # pylint: disable=protected-access
@@ -747,14 +798,14 @@ def __init__(
       self,
       flag_values=FLAGS,
       *,
-      accept_new_attributes: bool = False,
+      override_mode: OverrideMode = OverrideMode.NEVER,
       sys_argv=None,
       **kwargs,
   ):
     # Parent constructor can already call .Parse, thus additional fields
     # have to be set here.
     self.flag_values = flag_values
-    self._accept_new_attributes = accept_new_attributes
+    self._override_mode = override_mode
     # Note, we don't replace sys_argv with sys.argv here if it's None because
     # in some obscure multiprocessing use cases, sys.argv may not be populated
     # until later and we need to look it up at parse time.
@@ -839,7 +890,7 @@ def _parse(self, argument):
     self._initialize_missing_parent_fields(config, overrides)
     self._validate_overrides(config, overrides)
 
-    if self._accept_new_attributes:
+    if self._override_mode != OverrideMode.NEVER:
       # If user provide a new attribute, fallback to `object` to accept all
       # literal
       default_type = object
@@ -855,7 +906,9 @@ def _parse(self, argument):
       field_name = '{}.{}'.format(self.name, field_path)
 
       parser = None
-      if field_type in _FIELD_TYPE_TO_PARSER:
+      if self._override_mode == OverrideMode.ALWAYS:
+        parser = _LiteralParser()
+      elif field_type in _FIELD_TYPE_TO_PARSER:
         parser = _FIELD_TYPE_TO_PARSER[field_type]
       elif isinstance(field_type, type) and issubclass(
           field_type, config_dict.ConfigDict
@@ -903,7 +956,7 @@ def _parse(self, argument):
               serializer=serializer,
               name=field_name,
               default=default,
-              accept_new_attributes=self._accept_new_attributes,
+              override_mode=self._override_mode,
               help_string=field_help,
           )
           # Literal values support the `--my_bool` / `--nomy_bool` syntax
@@ -1073,7 +1126,7 @@ def __init__(
       help_string: str,
       short_name: Optional[str] = None,
       boolean: bool = False,
-      accept_new_attributes: bool = False,
+      override_mode: OverrideMode = OverrideMode.NEVER,
   ):
     """Creates new flag with callback."""
     super().__init__(
@@ -1087,15 +1140,23 @@ def __init__(
     self._path = path
     self._config = config
     self._override_values = override_values
-    self._accept_new_attributes = accept_new_attributes
+    self._override_mode = override_mode
 
   def parse(self, argument):
     super().parse(argument)
     # Callback to set value in ConfigDict.
-    config_path.set_value(
-        self._path, self._config, self.value,
-        accept_new_attributes=self._accept_new_attributes,
+    ctx = (
+        self._config.ignore_type()
+        if self._override_mode == OverrideMode.ALWAYS
+        else contextlib.nullcontext()
     )
+    with ctx:
+      config_path.set_value(
+          self._path,
+          self._config,
+          self.value,
+          accept_new_attributes=self._override_mode != OverrideMode.NEVER,
+      )
     self._override_values[self._path] = self.value
 
 

ml_collections/config_flags/tests/config_overriding_test.py

@@ -60,7 +60,8 @@ def _parse_flags(
     config=None,
     lock_config=True,
     required=False,
-    accept_new_attributes=False,
+    accept_new_attributes=None,
+    override_mode=None,
     use_sys_argv_override=False,
 ):
   """Parses arguments simulating sys.argv or via sys_argv argument."""
@@ -87,6 +88,7 @@ def _parse_flags(
         flag_values=values,
         lock_config=lock_config,
         accept_new_attributes=accept_new_attributes,
+        override_mode=override_mode,
         sys_argv=(argv if use_sys_argv_override else None),
     )
   else:
@@ -95,7 +97,10 @@ def _parse_flags(
         config=config,
         flag_values=values,
         lock_config=lock_config,
-        sys_argv=(argv if use_sys_argv_override else None))
+        accept_new_attributes=accept_new_attributes,
+        override_mode=override_mode,
+        sys_argv=(argv if use_sys_argv_override else None),
+    )
 
   if required:
     flags.mark_flag_as_required('test_config', flag_values=values)
@@ -726,6 +731,63 @@ def testNewAttributes(self):
     self.assertEqual(cfg.new_value, config_dict.ConfigDict({'a': [1, 2, 3]}))
     self.assertEqual(cfg.new_bool_value, True)
 
+  def testUseLiteralParser(self):
+    config = config_dict.ConfigDict({
+        'integer': 1,
+        'float': 2.0,
+        'string': 'hello',
+        'nested': {
+            'value': 10,
+        }
+    })
+    values = _parse_flags(
+        './program'
+        ' --test_config.integer="hello"'
+        ' --test_config.float=3'
+        ' --test_config.string="world"'
+        ' --test_config.nested.value=99'
+        ' --test_config.new_value="new"',
+        config=config,
+        lock_config=False,
+        override_mode=config_flags.OverrideMode.ALWAYS,
+    )
+    cfg = values.test_config
+    self.assertEqual(cfg.integer, 'hello')
+    self.assertEqual(cfg.float, 3)
+    self.assertEqual(cfg.string, 'world')
+    self.assertEqual(cfg.nested.value, 99)
+    self.assertEqual(cfg.new_value, 'new')
+
+  def testOverrideModeNewAttributes(self):
+    values = _parse_flags(
+        './program'
+        f' --test_config={_LITERAL_CONFIG_FILE}'
+        ' --test_config.integer=123'
+        ' --test_config.other_new_value="abc def"'
+        ' --test_config.new_value="{\'a\': [1, 2, 3]}"'
+        ' --test_config.new_bool_value=true',
+        override_mode=config_flags.OverrideMode.NEW_ONLY,
+        lock_config=False,
+    )
+    cfg = values.test_config
+    self.assertEqual(cfg.integer, 123)
+    self.assertIsNone(cfg.string)
+    self.assertEqual(cfg.other_new_value, 'abc def')
+    self.assertEqual(cfg.new_value, config_dict.ConfigDict({'a': [1, 2, 3]}))
+    self.assertEqual(cfg.new_bool_value, True)
+
+  def testOverrideModeBothParametersRaisesError(self):
+    with self.assertRaisesRegex(
+        ValueError,
+        'Cannot specify both `accept_new_attributes` and `override_mode`.',
+    ):
+      _parse_flags(
+          f'./program --test_config={_LITERAL_CONFIG_FILE}',
+          accept_new_attributes=True,
+          override_mode=config_flags.OverrideMode.NEW_ONLY,
+          lock_config=False,
+      )
+
 
 def _simple_config():
   config = config_dict.ConfigDict()