Internal change only.

PiperOrigin-RevId: 797749449

Author: GeorgOstrovski

ml_collections/config_dict/__init__.py

@@ -27,6 +27,11 @@
 from .config_dict import required_placeholder
 from .config_dict import RequiredValueError
 
+from .internal_utils import BestEffortCustomJSONEncoder
+from .internal_utils import Jsonable
+from .internal_utils import update_in_place
+from .internal_utils import copy_and_update
+
 __all__ = ("_Op", "ConfigDict", "create", "CustomJSONEncoder", "FieldReference",
            "FrozenConfigDict", "JSONDecodeError", "MutabilityError",
            "placeholder", "recursive_rename", "required_placeholder",

ml_collections/config_dict/config_dict.py

@@ -2108,3 +2108,84 @@ def recursive_rename(conf, old_name, new_name):
     else:
       setattr(new_conf, name, new_c)
   return new_conf
+
+
+# BEGIN GOOGLE-INTERNAL
+# pylint: disable=g-import-not-at-top,g-bad-import-order
+import copy
+# pylint: enable=g-import-not-at-top,g-bad-import-order
+
+
+def _setattr(
+    obj: Any,
+    key: str,
+    value: Any,
+    ensure_keys_exist: bool,
+    prefix: str,
+) -> None:
+  """Recursively setattr() for an object with a dot-separated key."""
+  if '.' in key:
+    key_first, key_rest = key.split('.', 1)
+    if not hasattr(obj, key_first):
+      if ensure_keys_exist:
+        full_key = prefix + key
+        child_key = prefix + key_first
+        raise AttributeError('Key "{}" cannot be set as "{}" was not found.'
+                             .format(full_key, child_key))
+      else:
+        setattr(obj, key_first, {})
+    _setattr(getattr(obj, key_first), key_rest, value,
+             ensure_keys_exist=ensure_keys_exist,
+             prefix=prefix + key_first + '.')
+  else:
+    if ensure_keys_exist and not hasattr(obj, key):
+      full_key = prefix + key
+      raise AttributeError(
+          f'Key "{full_key}" cannot be set as "{full_key}" was not found.')
+    setattr(obj, key, value)
+
+
+def update_in_place(
+    config: ConfigDict,
+    params: Mapping[str, Any],
+    ensure_keys_exist: bool = False,
+) -> ConfigDict:
+  """Returns the config dict updated in-place using dot-separated paths.
+
+  Args:
+    config: The config dict to update.
+    params: A dictionary containing the updated items. Keys are
+      dot-separated paths, and values are the updated values.
+    ensure_keys_exist: If True, raise an AttributeError, if a key does not
+      exist yet.
+
+  Returns:
+    The in-place updated config.
+  """
+  for key, value in params.items():
+    _setattr(config, key, value,
+             ensure_keys_exist=ensure_keys_exist,
+             prefix='')
+  return config
+
+
+def copy_and_update(
+    config: ConfigDict,
+    params: Mapping[str, Any],
+    ensure_keys_exist: bool = False,
+) -> ConfigDict:
+  """Copies a config dict and update using dot-separated paths.
+
+  Args:
+    config: The config dict to copy and update.
+    params: A dictionary containing the updated items. Keys are
+      dot-separated paths, and values are the updated values.
+    ensure_keys_exist: If True, raise an AttributeError, if a key does not exist
+      yet.
+
+  Returns:
+    A copy of the updated config dict.
+  """
+  new_config = copy.deepcopy(config)
+  return update_in_place(new_config, params, ensure_keys_exist)
+# END GOOGLE-INTERNAL