pyglove 0.4.5.dev202411132359__py3-none-any.whl → 0.4.5.dev202501250807__py3-none-any.whl

pyglove/core/typing/annotation_conversion.py

@@ -18,7 +18,7 @@ import inspect
 import types
 import typing
 
-from pyglove.core import object_utils
+from pyglove.core import utils
 from pyglove.core.typing import annotated
 from pyglove.core.typing import class_schema
 from pyglove.core.typing import inspect as pg_inspect
@@ -73,7 +73,7 @@ def _value_spec_from_default_value(
   elif isinstance(value, tuple):
     value_spec = vs.Tuple(
         [_value_spec_from_default_value(elem, False) for elem in value])
-  elif inspect.isfunction(value) or isinstance(value, object_utils.Functor):
+  elif inspect.isfunction(value) or isinstance(value, utils.Functor):
     value_spec = vs.Callable()
   elif not isinstance(value, type):
     value_spec = vs.Object(type(value))
@@ -132,7 +132,7 @@ def _value_spec_from_type_annotation(
     return vs.Union([vs.List(elem), vs.Tuple(elem)])
   # Handling literals.
   elif origin is typing.Literal:
-    return vs.Enum(object_utils.MISSING_VALUE, args)
+    return vs.Enum(utils.MISSING_VALUE, args)
   # Handling dict.
   elif origin in (dict, typing.Dict, collections.abc.Mapping):
     if not args:
@@ -177,6 +177,11 @@ def _value_spec_from_type_annotation(
     if optional:
       spec = spec.noneable()
     return spec
+  elif origin is typing.Final:
+    return _value_spec_from_type_annotation(
+        args[0],
+        accept_value_as_annotation=False
+    ).freeze(vs._FROZEN_VALUE_PLACEHOLDER)  # pylint: disable=protected-access
   elif isinstance(annotation, typing.ForwardRef):
     annotation = annotation.__forward_arg__
     if parent_module is not None:

pyglove/core/typing/annotation_conversion_test.py

@@ -306,6 +306,14 @@ class ValueSpecFromAnnotationTest(unittest.TestCase):
           ValueSpec.from_annotation(int | str, True),
           vs.Union([vs.Int(), vs.Str()]))
 
+  def test_final(self):
+    self.assertEqual(
+        ValueSpec.from_annotation(
+            typing.Final[int], True
+        ).set_default(1),
+        vs.Int().freeze(1)
+    )
+
 
 if __name__ == '__main__':
   unittest.main()

pyglove/core/typing/callable_ext.py

@@ -19,14 +19,14 @@ import inspect
 import types
 from typing import Any, Callable, Dict, Iterator, List, Optional, Sequence, Tuple, Union
 
-from pyglove.core import object_utils
+from pyglove.core import utils
 from pyglove.core.typing import callable_signature
 
 
 _TLS_KEY_PRESET_KWARGS = '__preset_kwargs__'
 
 
-class PresetArgValue(object_utils.Formattable):
+class PresetArgValue(utils.Formattable):
   """Value placeholder for arguments whose value will be provided by presets.
 
   Example:
@@ -39,12 +39,12 @@ class PresetArgValue(object_utils.Formattable):
     print(foo(x=1))  # 2: y=1
   """
 
-  def __init__(self, default: Any = object_utils.MISSING_VALUE):
+  def __init__(self, default: Any = utils.MISSING_VALUE):
     self.default = default
 
   @property
   def has_default(self) -> bool:
-    return self.default != object_utils.MISSING_VALUE
+    return self.default != utils.MISSING_VALUE
 
   def __eq__(self, other: Any) -> bool:
     return isinstance(other, PresetArgValue) and (
@@ -55,13 +55,13 @@ class PresetArgValue(object_utils.Formattable):
     return not self.__eq__(other)
 
   def format(self, *args, **kwargs):
-    return object_utils.kvlist_str(
+    return utils.kvlist_str(
         [
-            ('default', self.default, object_utils.MISSING_VALUE),
+            ('default', self.default, utils.MISSING_VALUE),
         ],
         label='PresetArgValue',
         *args,
-        **kwargs
+        **kwargs,
     )
 
   @classmethod
@@ -182,15 +182,15 @@ def preset_args(
     Current preset kwargs.
   """
 
-  parent_presets = object_utils.thread_local_peek(
+  parent_presets = utils.thread_local_peek(
       _TLS_KEY_PRESET_KWARGS, _ArgPresets()
   )
   current_preset = parent_presets.derive(kwargs, preset_name, inherit_preset)
-  object_utils.thread_local_push(_TLS_KEY_PRESET_KWARGS, current_preset)
+  utils.thread_local_push(_TLS_KEY_PRESET_KWARGS, current_preset)
   try:
     yield current_preset
   finally:
-    object_utils.thread_local_pop(_TLS_KEY_PRESET_KWARGS, None)
+    utils.thread_local_pop(_TLS_KEY_PRESET_KWARGS, None)
 
 
 def enable_preset_args(
@@ -243,9 +243,7 @@ def enable_preset_args(
       @functools.wraps(func)
       def _func(*args, **kwargs):
         # Map positional arguments to keyword arguments.
-        presets = object_utils.thread_local_peek(
-            _TLS_KEY_PRESET_KWARGS, None
-        )
+        presets = utils.thread_local_peek(_TLS_KEY_PRESET_KWARGS, None)
         preset_kwargs = presets.get_preset(preset_name) if presets else {}
         args, kwargs = PresetArgValue.resolve_args(
             args, kwargs, positional_arg_names, arg_defaults, preset_kwargs,

pyglove/core/typing/callable_signature.py

@@ -21,7 +21,8 @@ import types
 import typing
 from typing import Any, Callable, Dict, List, Optional, Union
 
-from pyglove.core import object_utils
+from pyglove.core import coding
+from pyglove.core import utils
 from pyglove.core.typing import class_schema
 from pyglove.core.typing import key_specs as ks
 
@@ -140,7 +141,7 @@ class CallableType(enum.Enum):
   METHOD = 2
 
 
-class Signature(object_utils.Formattable):
+class Signature(utils.Formattable):
   """PY3 function signature."""
 
   def __init__(self,
@@ -256,7 +257,7 @@ class Signature(object_utils.Formattable):
       **kwargs,
   ) -> str:
     """Format current object."""
-    return object_utils.kvlist_str(
+    return utils.kvlist_str(
         [
             ('', self.id, ''),
             ('args', self.args, []),
@@ -270,7 +271,7 @@ class Signature(object_utils.Formattable):
         compact=compact,
         verbose=verbose,
         root_indent=root_indent,
-        **kwargs
+        **kwargs,
     )
 
   def annotate(
@@ -287,7 +288,7 @@ class Signature(object_utils.Formattable):
       return_value = class_schema.ValueSpec.from_annotation(
           return_value, auto_typing=True
       )
-      if object_utils.MISSING_VALUE != return_value.default:
+      if utils.MISSING_VALUE != return_value.default:
         raise ValueError('return value spec should not have default value.')
       self.return_value = return_value
 
@@ -336,12 +337,12 @@ class Signature(object_utils.Formattable):
           or field.value.default is None
       ):
         field.value.set_default(
-            arg.value_spec.default, root_path=object_utils.KeyPath(arg.name)
+            arg.value_spec.default, root_path=utils.KeyPath(arg.name)
         )
       if arg.value_spec.default != field.value.default:
         if field.value.is_noneable and not arg.value_spec.has_default:
-            # Special handling noneable which always comes with a default.
-          field.value.set_default(object_utils.MISSING_VALUE)
+          # Special handling noneable which always comes with a default.
+          field.value.set_default(utils.MISSING_VALUE)
         elif not (
             # Special handling Dict type which always has default.
             isinstance(field.value, class_schema.ValueSpec.DictType)
@@ -548,7 +549,7 @@ class Signature(object_utils.Formattable):
     if not callable(callable_object):
       raise TypeError(f'{callable_object!r} is not callable.')
 
-    if isinstance(callable_object, object_utils.Functor):
+    if isinstance(callable_object, utils.Functor):
       assert callable_object.__signature__ is not None
       return callable_object.__signature__
 
@@ -565,15 +566,15 @@ class Signature(object_utils.Formattable):
         description = None
         args_doc = {}
         if func.__doc__:
-          cls_doc = object_utils.DocStr.parse(func.__doc__)
+          cls_doc = utils.DocStr.parse(func.__doc__)
           description = cls_doc.short_description
           args_doc.update(cls_doc.args)
 
         if func.__init__.__doc__:
-          init_doc = object_utils.DocStr.parse(func.__init__.__doc__)
+          init_doc = utils.DocStr.parse(func.__init__.__doc__)
           args_doc.update(init_doc.args)
-        docstr = object_utils.DocStr(
-            object_utils.DocStrStyle.GOOGLE,
+        docstr = utils.DocStr(
+            utils.DocStrStyle.GOOGLE,
             short_description=description,
             long_description=None,
             examples=[],
@@ -592,7 +593,7 @@ class Signature(object_utils.Formattable):
           else CallableType.FUNCTION
       )
       if auto_doc:
-        docstr = object_utils.docstr(func)
+        docstr = utils.docstr(func)
       sig = inspect.signature(func)
 
     module_name = getattr(func, '__module__', None)
@@ -616,7 +617,7 @@ class Signature(object_utils.Formattable):
       module_name: Optional[str] = None,
       qualname: Optional[str] = None,
       auto_typing: bool = False,
-      docstr: Union[str, object_utils.DocStr, None] = None,
+      docstr: Union[str, utils.DocStr, None] = None,
       parent_module: Optional[types.ModuleType] = None,
   ) -> 'Signature':
     """Returns PyGlove signature from Python signature.
@@ -643,7 +644,7 @@ class Signature(object_utils.Formattable):
     varkw = None
 
     if isinstance(docstr, str):
-      docstr = object_utils.DocStr.parse(docstr)
+      docstr = utils.DocStr.parse(docstr)
 
     def make_arg_spec(param: inspect.Parameter) -> Argument:
       """Makes argument spec from inspect.Parameter."""
@@ -707,7 +708,7 @@ class Signature(object_utils.Formattable):
         force_missing_as_default: bool = False,
         arg_prefix: str = ''):
       s = [f'{arg_prefix}{arg_name}']
-      if arg_spec.annotation != object_utils.MISSING_VALUE:
+      if arg_spec.annotation != utils.MISSING_VALUE:
         s.append(f': _annotation_{arg_name}')
         exec_locals[f'_annotation_{arg_name}'] = arg_spec.annotation
       if not arg_prefix and (force_missing_as_default or arg_spec.has_default):
@@ -753,14 +754,16 @@ class Signature(object_utils.Formattable):
       )
 
     # Generate function.
-    fn = object_utils.make_function(
+    fn = coding.make_function(
         self.name,
         args=args,
         body=body,
         exec_globals=exec_globals,
         exec_locals=exec_locals,
         return_type=getattr(
-            self.return_value, 'annotation', object_utils.MISSING_VALUE))
+            self.return_value, 'annotation', coding.NO_TYPE_ANNOTATION
+        ),
+    )
     fn.__module__ = self.module_name
     fn.__name__ = self.name
     fn.__qualname__ = self.qualname

pyglove/core/typing/callable_signature_test.py

@@ -11,15 +11,13 @@
 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 # See the License for the specific language governing permissions and
 # limitations under the License.
-"""Tests for pyglove.core.typing.callable_signature."""
-
 import copy
 import dataclasses
 import inspect
 from typing import List
 import unittest
 
-from pyglove.core import object_utils
+from pyglove.core import utils
 from pyglove.core.typing import annotation_conversion   # pylint: disable=unused-import
 from pyglove.core.typing import callable_signature
 from pyglove.core.typing import class_schema
@@ -628,8 +626,8 @@ class FromCallableTest(unittest.TestCase):
     self.assertIsNotNone(signature.varkw)
 
   def test_signature_with_forward_declarations(self):
-    signature = callable_signature.signature(object_utils.KeyPath)
-    self.assertIs(signature.get_value_spec('parent').cls, object_utils.KeyPath)
+    signature = callable_signature.signature(utils.KeyPath)
+    self.assertIs(signature.get_value_spec('parent').cls, utils.KeyPath)
 
 
 class FromSchemaTest(unittest.TestCase):

pyglove/core/typing/class_schema.py

@@ -20,10 +20,10 @@ import sys
 import types
 from typing import Any, Callable, Dict, Iterable, List, Optional, Sequence, Set, Tuple, Type, Union
 
-from pyglove.core import object_utils
+from pyglove.core import utils
 
 
-class KeySpec(object_utils.Formattable, object_utils.JSONConvertible):
+class KeySpec(utils.Formattable, utils.JSONConvertible):
   """Interface for key specifications.
 
   A key specification determines what keys are acceptable for a symbolic
@@ -94,7 +94,7 @@ class KeySpec(object_utils.Formattable, object_utils.JSONConvertible):
     assert False, 'Overridden in `key_specs.py`.'
 
 
-class ForwardRef(object_utils.Formattable):
+class ForwardRef(utils.Formattable):
   """Forward type reference."""
 
   def __init__(self, module: types.ModuleType, name: str):
@@ -147,7 +147,7 @@ class ForwardRef(object_utils.Formattable):
       **kwargs
   ) -> str:
     """Format this object."""
-    return object_utils.kvlist_str(
+    return utils.kvlist_str(
         [
             ('module', self.module.__name__, None),
             ('name', self.name, None),
@@ -180,7 +180,7 @@ class ForwardRef(object_utils.Formattable):
     return ForwardRef(self.module, self.name)
 
 
-class ValueSpec(object_utils.Formattable, object_utils.JSONConvertible):
+class ValueSpec(utils.Formattable, utils.JSONConvertible):
   """Interface for value specifications.
 
   A value specification defines what values are acceptable for a symbolic
@@ -344,6 +344,10 @@ class ValueSpec(object_utils.Formattable, object_utils.JSONConvertible):
       Tuple[Type[Any], ...]]:  # pyformat: disable
     """Returns acceptable (resolved) value type(s)."""
 
+  @abc.abstractmethod
+  def __call__(self, *args, **kwargs) -> Any:
+    """Instantiates a value based on the spec.."""
+
   @property
   @abc.abstractmethod
   def forward_refs(self) -> Set[ForwardRef]:
@@ -363,7 +367,7 @@ class ValueSpec(object_utils.Formattable, object_utils.JSONConvertible):
       self,
       default: Any,
       use_default_apply: bool = True,
-      root_path: Optional[object_utils.KeyPath] = None
+      root_path: Optional[utils.KeyPath] = None,
   ) -> 'ValueSpec':
     """Sets the default value and returns `self`.
 
@@ -394,13 +398,14 @@ class ValueSpec(object_utils.Formattable, object_utils.JSONConvertible):
   @property
   def has_default(self) -> bool:
     """Returns True if the default value is provided."""
-    return self.default != object_utils.MISSING_VALUE
+    return self.default != utils.MISSING_VALUE
 
   @abc.abstractmethod
   def freeze(
       self,
-      permanent_value: Any = object_utils.MISSING_VALUE,
-      apply_before_use: bool = True) -> 'ValueSpec':
+      permanent_value: Any = utils.MISSING_VALUE,
+      apply_before_use: bool = True,
+  ) -> 'ValueSpec':
     """Sets the default value using a permanent value and freezes current spec.
 
     A frozen value spec will not accept any value that is not the default
@@ -467,10 +472,11 @@ class ValueSpec(object_utils.Formattable, object_utils.JSONConvertible):
       self,
       value: Any,
       allow_partial: bool = False,
-      child_transform: Optional[Callable[
-          [object_utils.KeyPath, 'Field', Any], Any]] = None,
-      root_path: Optional[object_utils.KeyPath] = None,
-      ) -> Any:
+      child_transform: Optional[
+          Callable[[utils.KeyPath, 'Field', Any], Any]
+      ] = None,
+      root_path: Optional[utils.KeyPath] = None,
+  ) -> Any:
     """Validates, completes and transforms the input value.
 
     Here is the procedure of ``apply``::
@@ -547,7 +553,7 @@ class ValueSpec(object_utils.Formattable, object_utils.JSONConvertible):
     assert False, 'Overridden in `annotation_conversion.py`.'
 
 
-class Field(object_utils.Formattable, object_utils.JSONConvertible):
+class Field(utils.Formattable, utils.JSONConvertible):
   """Class that represents the definition of one or a group of attributes.
 
   ``Field`` is held by a :class:`pyglove.Schema` object for defining the
@@ -587,7 +593,9 @@ class Field(object_utils.Formattable, object_utils.JSONConvertible):
       key_spec: Union[KeySpec, str],
       value_spec: ValueSpec,
       description: Optional[str] = None,
-      metadata: Optional[Dict[str, Any]] = None):
+      metadata: Optional[Dict[str, Any]] = None,
+      origin: Optional[Type[Any]] = None,
+  ) -> None:
     """Constructor.
 
     Args:
@@ -596,6 +604,7 @@ class Field(object_utils.Formattable, object_utils.JSONConvertible):
       value_spec: Value specification of the field.
       description: Description of the field.
       metadata: A dict of objects as metadata for the field.
+      origin: The class that this field originates from.
 
     Raises:
       ValueError: metadata is not a dict.
@@ -606,6 +615,7 @@ class Field(object_utils.Formattable, object_utils.JSONConvertible):
     self._key = key_spec
     self._value = value_spec
     self._description = description
+    self._origin = origin
 
     if metadata and not isinstance(metadata, dict):
       raise ValueError('metadata must be a dict.')
@@ -661,6 +671,15 @@ class Field(object_utils.Formattable, object_utils.JSONConvertible):
     """
     return self._metadata
 
+  @property
+  def origin(self) -> Optional[Type[Any]]:
+    """The class that this field originates from."""
+    return self._origin
+
+  def set_origin(self, origin: Type[Any]) -> None:
+    """Sets the origin (source class) of this field."""
+    self._origin = origin
+
   def extend(self, base_field: 'Field') -> 'Field':
     """Extend current field based on a base field."""
     self.key.extend(base_field.key)
@@ -677,9 +696,11 @@ class Field(object_utils.Formattable, object_utils.JSONConvertible):
       self,
       value: Any,
       allow_partial: bool = False,
-      transform_fn: Optional[Callable[
-          [object_utils.KeyPath, 'Field', Any], Any]] = None,
-      root_path: Optional[object_utils.KeyPath] = None) -> Any:
+      transform_fn: Optional[
+          Callable[[utils.KeyPath, 'Field', Any], Any]
+      ] = None,
+      root_path: Optional[utils.KeyPath] = None,
+  ) -> Any:
     """Apply current field to a value, which validate and complete the value.
 
     Args:
@@ -731,18 +752,19 @@ class Field(object_utils.Formattable, object_utils.JSONConvertible):
       **kwargs,
   ) -> str:
     """Format this field into a string."""
-    return object_utils.kvlist_str(
+    return utils.kvlist_str(
         [
             ('key', self._key, None),
             ('value', self._value, None),
             ('description', self._description, None),
             ('metadata', self._metadata, {}),
+            ('origin', self._origin, None),
         ],
         label=self.__class__.__name__,
         compact=compact,
         verbose=verbose,
         root_indent=root_indent,
-        **kwargs
+        **kwargs,
     )
 
   def to_json(self, **kwargs: Any) -> Dict[str, Any]:
@@ -771,7 +793,7 @@ class Field(object_utils.Formattable, object_utils.JSONConvertible):
     return not self.__eq__(other)
 
 
-class Schema(object_utils.Formattable, object_utils.JSONConvertible):
+class Schema(utils.Formattable, utils.JSONConvertible):
   """Class that represents a schema.
 
   PyGlove's runtime type system is based on the concept of ``Schema`` (
@@ -861,7 +883,9 @@ class Schema(object_utils.Formattable, object_utils.JSONConvertible):
       description: Optional[str] = None,
       *,
       allow_nonconst_keys: bool = False,
-      metadata: Optional[Dict[str, Any]] = None):
+      metadata: Optional[Dict[str, Any]] = None,
+      for_cls: Optional[Type[Any]] = None,
+  ):
     """Constructor.
 
     Args:
@@ -874,6 +898,7 @@ class Schema(object_utils.Formattable, object_utils.JSONConvertible):
       description: Optional str as the description for the schema.
       allow_nonconst_keys: Whether immediate fields can use non-const keys.
       metadata: Optional dict of user objects as schema-level metadata.
+      for_cls: Optional class that this schema applies to.
 
     Raises:
       TypeError: Argument `fields` is not a list.
@@ -895,6 +920,11 @@ class Schema(object_utils.Formattable, object_utils.JSONConvertible):
     self._description = description
     self._metadata = metadata or {}
 
+    if for_cls is not None:
+      for f in fields:
+        if f.origin is None:
+          f.set_origin(for_cls)
+
     self._dynamic_field = None
     for f in fields:
       if not f.key.is_const:
@@ -930,21 +960,27 @@ class Schema(object_utils.Formattable, object_utils.JSONConvertible):
     Returns:
       The merged schema.
     """
-    field_names = set()
-    fields = []
+    fields = {}
     kw_field = None
     for schema in schema_list:
       for key, field in schema.fields.items():
-        if key.is_const and key not in field_names:
-          fields.append(field)
-          field_names.add(key)
-        elif not key.is_const and kw_field is None:
+        if key.is_const:
+          if key not in fields or (
+              field.origin is not None
+              and fields[key].origin is not None
+              and issubclass(field.origin, fields[key].origin)
+          ):
+            fields[key] = field
+        elif kw_field is None:
           kw_field = field
 
     if kw_field is not None:
-      fields.append(kw_field)
+      fields[kw_field.key] = kw_field
     return Schema(
-        fields, name=name, description=description, allow_nonconst_keys=True
+        list(fields.values()),
+        name=name,
+        description=description,
+        allow_nonconst_keys=True
     )
 
   def extend(self, base: 'Schema') -> 'Schema':
@@ -955,13 +991,12 @@ class Schema(object_utils.Formattable, object_utils.JSONConvertible):
         parent_field: Field,
         child_field: Field) -> Field:
       """Merge function on field with the same key."""
-      if parent_field != object_utils.MISSING_VALUE:
-        if object_utils.MISSING_VALUE == child_field:
+      if parent_field != utils.MISSING_VALUE:
+        if utils.MISSING_VALUE == child_field:
           if (not self._allow_nonconst_keys and not parent_field.key.is_const):
-            hints = object_utils.kvlist_str([
-                ('base', base.name, None),
-                ('path', path, None)
-            ])
+            hints = utils.kvlist_str(
+                [('base', base.name, None), ('path', path, None)]
+            )
             raise ValueError(
                 f'Non-const key {parent_field.key} is not allowed to be '
                 f'added to the schema. ({hints})')
@@ -970,16 +1005,15 @@ class Schema(object_utils.Formattable, object_utils.JSONConvertible):
           try:
             child_field.extend(parent_field)
           except Exception as e:  # pylint: disable=broad-except
-            hints = object_utils.kvlist_str([
-                ('base', base.name, None),
-                ('path', path, None)
-            ])
+            hints = utils.kvlist_str(
+                [('base', base.name, None), ('path', path, None)]
+            )
             raise e.__class__(f'{e} ({hints})').with_traceback(
                 sys.exc_info()[2])
       return child_field
 
-    self._fields = object_utils.merge([base.fields, self.fields], _merge_field)
-    self._metadata = object_utils.merge([base.metadata, self.metadata])
+    self._fields = utils.merge([base.fields, self.fields], _merge_field)
+    self._metadata = utils.merge([base.metadata, self.metadata])
 
     # Inherit dynamic field from base if it's not present in the child.
     if self._dynamic_field is None:
@@ -1102,8 +1136,8 @@ class Schema(object_utils.Formattable, object_utils.JSONConvertible):
       dict_obj: Dict[str, Any],
       allow_partial: bool = False,
       child_transform: Optional[Callable[
-          [object_utils.KeyPath, Field, Any], Any]] = None,
-      root_path: Optional[object_utils.KeyPath] = None,
+          [utils.KeyPath, Field, Any], Any]] = None,
+      root_path: Optional[utils.KeyPath] = None,
   ) -> Dict[str, Any]:  # pyformat: disable
     # pyformat: disable
     """Apply this schema to a dict object, validate and transform it.
@@ -1160,18 +1194,18 @@ class Schema(object_utils.Formattable, object_utils.JSONConvertible):
         keys.append(str(key_spec))
       for key in keys:
         if dict_obj:
-          value = dict_obj.get(key, object_utils.MISSING_VALUE)
+          value = dict_obj.get(key, utils.MISSING_VALUE)
         else:
-          value = object_utils.MISSING_VALUE
+          value = utils.MISSING_VALUE
         # NOTE(daiyip): field.default_value may be MISSING_VALUE too
         # or partial.
-        if object_utils.MISSING_VALUE == value:
+        if utils.MISSING_VALUE == value:
           value = copy.deepcopy(field.default_value)
         new_value = field.apply(
             value,
             allow_partial=allow_partial,
             transform_fn=child_transform,
-            root_path=object_utils.KeyPath(key, root_path)
+            root_path=utils.KeyPath(key, root_path),
         )
 
         # NOTE(daiyip): `pg.Dict.__getitem__`` has special logics in handling
@@ -1185,10 +1219,12 @@ class Schema(object_utils.Formattable, object_utils.JSONConvertible):
           dict_obj[key] = new_value
     return dict_obj
 
-  def validate(self,
-               dict_obj: Dict[str, Any],
-               allow_partial: bool = False,
-               root_path: Optional[object_utils.KeyPath] = None) -> None:
+  def validate(
+      self,
+      dict_obj: Dict[str, Any],
+      allow_partial: bool = False,
+      root_path: Optional[utils.KeyPath] = None,
+  ) -> None:
     """Validates whether dict object is conformed with the schema."""
     self.apply(
         copy.deepcopy(dict_obj),
@@ -1253,12 +1289,12 @@ class Schema(object_utils.Formattable, object_utils.JSONConvertible):
       root_indent: int = 0,
       *,
       cls_name: Optional[str] = None,
-      bracket_type: object_utils.BracketType = object_utils.BracketType.ROUND,
+      bracket_type: utils.BracketType = utils.BracketType.ROUND,
       fields_only: bool = False,
       **kwargs,
   ) -> str:
     """Format current Schema into nicely printed string."""
-    return object_utils.kvlist_str(
+    return utils.kvlist_str(
         [
             ('name', self.name, None),
             ('description', self.description, None),
@@ -1407,6 +1443,7 @@ def create_schema(
     allow_nonconst_keys: bool = False,
     metadata: Optional[Dict[str, Any]] = None,
     description: Optional[str] = None,
+    for_cls: Optional[Type[Any]] = None,
     parent_module: Optional[types.ModuleType] = None
 ) -> Schema:
   """Creates ``Schema`` from a list of ``Field``s or equivalences.
@@ -1448,6 +1485,7 @@ def create_schema(
     allow_nonconst_keys: Whether to allow non const keys in schema.
     metadata: Optional dict of user objects as schema-level metadata.
     description: Optional description of the schema.
+    for_cls: (Optional) the class that this schema applies to.
     parent_module: (Optional) parent module for defining this schema, which will
       be used for forward reference lookup.
 
@@ -1471,6 +1509,7 @@ def create_schema(
       allow_nonconst_keys=allow_nonconst_keys,
       metadata=metadata,
       description=description,
+      for_cls=for_cls,
   )