pyglove 0.4.5.dev20240318__py3-none-any.whl → 0.4.5.dev202501132210__py3-none-any.whl

pyglove/core/symbolic/base_test.py

@@ -11,15 +11,18 @@
 # 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.symbolic.base."""
-
 import copy
+import inspect
+from typing import Any
 import unittest
 
-from pyglove.core import object_utils
 from pyglove.core import typing as pg_typing
+from pyglove.core import utils
+from pyglove.core import views
 from pyglove.core.symbolic import base
 from pyglove.core.symbolic.dict import Dict
+from pyglove.core.symbolic.inferred import ValueFromParentChain
+from pyglove.core.symbolic.object import Object
 
 
 class FieldUpdateTest(unittest.TestCase):
@@ -28,7 +31,7 @@ class FieldUpdateTest(unittest.TestCase):
   def test_basics(self):
     x = Dict(x=1)
     f = pg_typing.Field('x', pg_typing.Int())
-    update = base.FieldUpdate(object_utils.KeyPath('x'), x, f, 1, 2)
+    update = base.FieldUpdate(utils.KeyPath('x'), x, f, 1, 2)
     self.assertEqual(update.path, 'x')
     self.assertIs(update.target, x)
     self.assertIs(update.field, f)
@@ -37,15 +40,15 @@ class FieldUpdateTest(unittest.TestCase):
 
   def test_format(self):
     self.assertEqual(
-        base.FieldUpdate(
-            object_utils.KeyPath('x'), Dict(x=1), None, 1, 2
-        ).format(compact=True),
+        base.FieldUpdate(utils.KeyPath('x'), Dict(x=1), None, 1, 2).format(
+            compact=True
+        ),
         'FieldUpdate(parent_path=, path=x, old_value=1, new_value=2)',
     )
 
     self.assertEqual(
         base.FieldUpdate(
-            object_utils.KeyPath('a'), Dict(x=Dict(a=1)).x, None, 1, 2
+            utils.KeyPath('a'), Dict(x=Dict(a=1)).x, None, 1, 2
         ).format(compact=True),
         'FieldUpdate(parent_path=x, path=a, old_value=1, new_value=2)',
     )
@@ -54,35 +57,136 @@ class FieldUpdateTest(unittest.TestCase):
     x = Dict()
     f = pg_typing.Field('x', pg_typing.Int())
     self.assertEqual(
-        base.FieldUpdate(object_utils.KeyPath('a'), x, f, 1, 2),
-        base.FieldUpdate(object_utils.KeyPath('a'), x, f, 1, 2),
+        base.FieldUpdate(utils.KeyPath('a'), x, f, 1, 2),
+        base.FieldUpdate(utils.KeyPath('a'), x, f, 1, 2),
     )
 
     # Targets are not the same instance.
     self.assertNotEqual(
-        base.FieldUpdate(object_utils.KeyPath('a'), x, f, 1, 2),
-        base.FieldUpdate(object_utils.KeyPath('a'), Dict(), f, 1, 2),
+        base.FieldUpdate(utils.KeyPath('a'), x, f, 1, 2),
+        base.FieldUpdate(utils.KeyPath('a'), Dict(), f, 1, 2),
     )
 
     # Fields are not the same instance.
     self.assertNotEqual(
-        base.FieldUpdate(object_utils.KeyPath('a'), x, f, 1, 2),
-        base.FieldUpdate(object_utils.KeyPath('b'), x, copy.copy(f), 1, 2),
+        base.FieldUpdate(utils.KeyPath('a'), x, f, 1, 2),
+        base.FieldUpdate(utils.KeyPath('b'), x, copy.copy(f), 1, 2),
     )
 
     self.assertNotEqual(
-        base.FieldUpdate(object_utils.KeyPath('a'), x, f, 1, 2),
-        base.FieldUpdate(object_utils.KeyPath('a'), x, f, 0, 2),
+        base.FieldUpdate(utils.KeyPath('a'), x, f, 1, 2),
+        base.FieldUpdate(utils.KeyPath('a'), x, f, 0, 2),
     )
 
     self.assertNotEqual(
-        base.FieldUpdate(object_utils.KeyPath('a'), x, f, 1, 2),
-        base.FieldUpdate(object_utils.KeyPath('a'), x, f, 1, 1),
+        base.FieldUpdate(utils.KeyPath('a'), x, f, 1, 2),
+        base.FieldUpdate(utils.KeyPath('a'), x, f, 1, 1),
     )
 
     self.assertNotEqual(
-        base.FieldUpdate(object_utils.KeyPath('a'), x, f, 1, 2), Dict()
+        base.FieldUpdate(utils.KeyPath('a'), x, f, 1, 2), Dict()
+    )
+
+
+class HtmlTreeViewExtensionTest(unittest.TestCase):
+
+  def assert_content(self, html, expected):
+    expected = inspect.cleandoc(expected).strip()
+    actual = html.content.strip()
+    if actual != expected:
+      print(actual)
+    self.assertEqual(actual.strip(), expected)
+
+  def test_to_html(self):
+
+    class Foo(Object):
+      x: int
+      y: Any = 'foo'
+      z: pg_typing.Int().freeze(1)
+
+    # Disable tooltip.
+    self.assert_content(
+        Foo(x=1, y='foo').to_html(
+            enable_summary_tooltip=False,
+            enable_key_tooltip=False
+        ),
+        """
+        <details open class="pyglove foo"><summary><div class="summary-title">Foo(...)</div></summary><div class="complex-value foo"><details open class="pyglove int"><summary><div class="summary-name">x</div><div class="summary-title">int</div></summary><span class="simple-value int">1</span></details><details open class="pyglove str"><summary><div class="summary-name">y</div><div class="summary-title">str</div></summary><span class="simple-value str">&#x27;foo&#x27;</span></details></div></details>
+        """
+    )
+    # Hide frozen and default values.
+    self.assert_content(
+        Foo(x=1, y='foo').to_html(
+            enable_summary_tooltip=False,
+            enable_key_tooltip=False,
+            collapse_level=0,
+            key_style='label',
+            extra_flags=dict(
+                hide_frozen=True,
+                hide_default_values=True
+            )
+        ),
+        """
+        <details class="pyglove foo"><summary><div class="summary-title">Foo(...)</div></summary><div class="complex-value foo"><table><tr><td><span class="object-key str">x</span></td><td><span class="simple-value int">1</span></td></tr></table></div></details>
+        """
+    )
+    self.assert_content(
+        Foo(x=1, y='foo').to_html(
+            enable_summary_tooltip=False,
+            enable_key_tooltip=False,
+            collapse_level=0,
+            extra_flags=dict(
+                hide_frozen=True,
+                hide_default_values=True
+            )
+        ),
+        """
+        <details class="pyglove foo"><summary><div class="summary-title">Foo(...)</div></summary><div class="complex-value foo"><details open class="pyglove int"><summary><div class="summary-name">x</div><div class="summary-title">int</div></summary><span class="simple-value int">1</span></details></div></details>
+        """
+    )
+    # Use inferred values.
+    x = Dict(x=Dict(y=ValueFromParentChain()), y=2)
+    self.assert_content(
+        x.x.to_html(
+            enable_summary_tooltip=False,
+            enable_key_tooltip=False,
+            key_style='label',
+            extra_flags=dict(
+                use_inferred=False
+            )
+        ),
+        """
+        <details open class="pyglove dict"><summary><div class="summary-title">Dict(...)</div></summary><div class="complex-value dict"><table><tr><td><span class="object-key str">y</span></td><td><details class="pyglove value-from-parent-chain"><summary><div class="summary-title">ValueFromParentChain(...)</div></summary><div class="complex-value value-from-parent-chain"><span class="empty-container"></span></div></details></td></tr></table></div></details>
+        """
+    )
+    self.assert_content(
+        x.x.to_html(
+            enable_summary_tooltip=False,
+            enable_key_tooltip=False,
+            key_style='label',
+            extra_flags=dict(
+                use_inferred=True
+            )
+        ),
+        """
+        <details open class="pyglove dict"><summary><div class="summary-title">Dict(...)</div></summary><div class="complex-value dict"><table><tr><td><span class="object-key str">y</span></td><td><span class="simple-value int">2</span></td></tr></table></div></details>
+        """
     )
+    # Test collapse level.
+    v = Foo(1, Foo(2, Foo(3, Foo(4))))
+    with views.view_options(key_style='label'):
+      self.assertEqual(
+          v.to_html(collapse_level=0).content.count('open'), 0
+      )
+      self.assertEqual(
+          v.to_html(collapse_level=1).content.count('open'), 1
+      )
+      self.assertEqual(
+          v.to_html(collapse_level=2).content.count('open'), 2
+      )
+      self.assertEqual(
+          v.to_html(collapse_level=None).content.count('open'), 4
+      )
 
 
 if __name__ == '__main__':

pyglove/core/symbolic/boilerplate.py

@@ -15,14 +15,12 @@
 
 import copy
 import inspect
-
 from typing import Any, List, Optional, Type
 
-from pyglove.core import object_utils
 from pyglove.core import typing as pg_typing
+from pyglove.core import utils
 from pyglove.core.symbolic import flags
 from pyglove.core.symbolic import object as pg_object
-from pyglove.core.symbolic import schema_utils
 
 
 def boilerplate_class(
@@ -130,9 +128,9 @@ def boilerplate_class(
   cls.auto_register = True
 
   allow_partial = value.allow_partial
-  def _freeze_field(path: object_utils.KeyPath,
-                    field: pg_typing.Field,
-                    value: Any) -> Any:
+  def _freeze_field(
+      path: utils.KeyPath, field: pg_typing.Field, value: Any
+  ) -> Any:
     # We do not do validation since Object is already in valid form.
     del path
     if not isinstance(field.key, pg_typing.ListKey):
@@ -152,18 +150,14 @@ def boilerplate_class(
     return value
 
   # NOTE(daiyip): we call `cls.__schema__.apply` to freeze fields that have
-  # default values. But we no longer need to formalize `cls.__schema__`, since
-  # it's copied from the boilerplate object's class which was already
-  # formalized.
+  # default values.
   with flags.allow_writable_accessors():
     cls.__schema__.apply(
         value._sym_attributes,  # pylint: disable=protected-access
         allow_partial=allow_partial,
         child_transform=_freeze_field,
     )
-
-  if init_arg_list is not None:
-    schema_utils.validate_init_arg_list(init_arg_list, cls.__schema__)
-    cls.__schema__.metadata['init_arg_list'] = init_arg_list
+  cls.__schema__.metadata['init_arg_list'] = init_arg_list
+  cls.apply_schema(cls.__schema__)
   cls.register_for_deserialization(serialization_key, additional_keys)
   return cls

pyglove/core/symbolic/boilerplate_test.py

@@ -70,35 +70,37 @@ class BoilerplateClassTest(unittest.TestCase):
       pg_boilerplate_class('A', template_object, init_arg_list=['x', 'y'])
 
   def test_init_arg_list(self):
-    self.assertEqual(B.init_arg_list, ['a', 'b', 'c'])
+    self.assertEqual(B.init_arg_list, ['a'])
     self.assertEqual(C.init_arg_list, ['a', 'c', 'b'])
 
   def test_schema(self):
     # Boilerplate class' schema should carry the default value and be frozen.
     self.assertEqual(
-        B.__schema__,
-        pg_typing.create_schema([
-            ('a', pg_typing.Int()),
-            (
-                'b',
-                pg_typing.Union(
-                    [pg_typing.Int(), pg_typing.Str()], default='foo'
-                ).freeze(),
-            ),
-            (
-                'c',
-                pg_typing.Dict([(
-                    'd',
-                    pg_typing.List(
-                        pg_typing.Dict([
-                            ('e', pg_typing.Float()),
-                            ('f', pg_typing.Bool()),
-                        ]),
-                        default=List([Dict(e=1.0, f=True)]),
+        list(B.__schema__.fields.values()),
+        list(pg_typing.create_schema(
+            [
+                ('a', pg_typing.Int()),
+                (
+                    'b',
+                    pg_typing.Union(
+                        [pg_typing.Int(), pg_typing.Str()], default='foo'
                     ).freeze(),
-                )]).freeze(),
-            ),
-        ]),
+                ),
+                (
+                    'c',
+                    pg_typing.Dict([(
+                        'd',
+                        pg_typing.List(
+                            pg_typing.Dict([
+                                ('e', pg_typing.Float()),
+                                ('f', pg_typing.Bool()),
+                            ]),
+                            default=List([Dict(e=1.0, f=True)]),
+                        ).freeze(),
+                    )]).freeze(),
+                ),
+            ],
+        ).fields.values())
     )
 
     # Original class' schema should remain unchanged.

pyglove/core/symbolic/class_wrapper.py

@@ -23,22 +23,21 @@ import abc
 import functools
 import inspect
 import types
-from typing import Any, Callable, Dict, List, Optional, Sequence, Text, Tuple, Type, Union
+from typing import Any, Callable, Dict, List, Optional, Sequence, Tuple, Type, Union
 
 from pyglove.core import detouring
-from pyglove.core import object_utils
 from pyglove.core import typing as pg_typing
+from pyglove.core import utils
 
 from pyglove.core.symbolic import dict as pg_dict  # pylint: disable=unused-import
 from pyglove.core.symbolic import list as pg_list  # pylint: disable=unused-import
 from pyglove.core.symbolic import object as pg_object
-from pyglove.core.symbolic import schema_utils
 
 
 class ClassWrapperMeta(pg_object.ObjectMeta):
   """Metaclass for class wrapper."""
 
-  def __repr__(self) -> Text:
+  def __repr__(self) -> str:
     wrapped_cls = getattr(self, 'sym_wrapped_cls', None)
     if wrapped_cls is None:
       return f'<class {self.__type_name__!r}>'
@@ -46,11 +45,7 @@ class ClassWrapperMeta(pg_object.ObjectMeta):
 
   def __getattr__(self, name):
     """Pass through attribute requests to sym_wrapped_cls."""
-    try:
-      return super().__getattr__(name)
-    except AttributeError:
-      wrapped_cls = object.__getattribute__(self, 'sym_wrapped_cls')
-      return getattr(wrapped_cls, name)
+    return getattr(object.__getattribute__(self, 'sym_wrapped_cls'), name)
 
 
 class ClassWrapper(pg_object.Object, metaclass=ClassWrapperMeta):
@@ -76,7 +71,7 @@ class _SubclassedWrapperBase(ClassWrapper):
   # the `__init__` method.
   auto_typing = False
 
-  @object_utils.explicit_method_override
+  @utils.explicit_method_override
   def __init__(self, *args, **kwargs):
     """Overridden __init__ to construct symbolic wrapper only."""
     # NOTE(daiyip): We avoid `__init__` to be called multiple times.
@@ -105,7 +100,7 @@ class _SubclassedWrapperBase(ClassWrapper):
   def __init_subclass__(cls):
     # Class wrappers inherit `__init__` from the user class. Therefore, we mark
     # all of them as explicitly overridden.
-    object_utils.explicit_method_override(cls.__init__)
+    utils.explicit_method_override(cls.__init__)
 
     super().__init_subclass__()
     if cls.__init__ is _SubclassedWrapperBase.__init__:
@@ -131,10 +126,10 @@ class _SubclassedWrapperBase(ClassWrapper):
       # In both cases, we need to generate an __init__ wrapper for
       # calling the symbolic initialization.
       setattr(cls, '__orig_init__', cls.__init__)
-      description, init_arg_list, arg_fields = _extract_init_signature(
+      init_arg_list, arg_fields = _extract_init_signature(
           cls, auto_doc=cls.auto_doc, auto_typing=cls.auto_typing)
 
-      @object_utils.explicit_method_override
+      @utils.explicit_method_override
       @functools.wraps(cls.__init__)
       def _sym_init(self, *args, **kwargs):
         _SubclassedWrapperBase.__init__(self, *args, **kwargs)
@@ -152,9 +147,7 @@ class _SubclassedWrapperBase(ClassWrapper):
 
       # We do not extend existing schema which is inherited from the base
       # class.
-      schema_utils.update_schema(
-          cls, arg_fields, init_arg_list=init_arg_list,
-          extend=False, description=description)
+      cls.update_schema(arg_fields, init_arg_list=init_arg_list, extend=False)
     else:
       assert hasattr(cls, '__orig_init__')
 
@@ -228,8 +221,8 @@ def _subclassed_wrapper(
     use_auto_doc: bool,
     use_auto_typing: bool,
     reset_state_fn: Optional[Callable[[Any], None]],
-    class_name: Optional[Text] = None,
-    module_name: Optional[Text] = None):
+    class_name: Optional[str] = None,
+    module_name: Optional[str] = None):
   """Class wrapper implementation by regular multi-inheritance."""
   # NOTE(daiyip): The user class may have a user-defined metaclass, which
   # conflicts with the metaclass of the symbolic base. Therefore, we detect
@@ -312,10 +305,11 @@ def _subclassed_wrapper(
 
 def wrap(
     cls,
-    init_args: Optional[List[Union[
-        Tuple[Union[Text, pg_typing.KeySpec], pg_typing.ValueSpec, Text],
-        Tuple[Union[Text, pg_typing.KeySpec], pg_typing.ValueSpec, Text, Any]
-    ]]] = None,
+    init_args: Union[
+        List[Union[pg_typing.Field, pg_typing.FieldDef]],
+        Dict[pg_typing.FieldKeyDef, pg_typing.FieldValueDef],
+        None
+    ] = None,
     *,
     reset_state_fn: Optional[Callable[[Any], None]] = None,
     repr: bool = True,    # pylint: disable=redefined-builtin
@@ -429,16 +423,10 @@ def wrap(
   if issubclass(cls, ClassWrapper):
     # Update init argument specifications according to user specified specs.
     # Replace schema instead of extending it.
-    description, init_arg_list, arg_fields = _extract_init_signature(
+    init_arg_list, arg_fields = _extract_init_signature(
         cls, init_args, auto_doc=auto_doc, auto_typing=auto_typing)
-    schema_utils.update_schema(
-        cls,
-        arg_fields,
-        init_arg_list=init_arg_list,
-        extend=False,
-        description=description,
-        serialization_key=serialization_key,
-        additional_keys=additional_keys)
+    cls.update_schema(arg_fields, init_arg_list=init_arg_list, extend=False)
+    cls.register_for_deserialization(serialization_key, additional_keys)
 
   if override:
     for k, v in override.items():
@@ -448,7 +436,7 @@ def wrap(
 
 def wrap_module(
     module,
-    names: Optional[Sequence[Text]] = None,
+    names: Optional[Sequence[str]] = None,
     where: Optional[Callable[[Type['ClassWrapper']], bool]] = None,
     export_to: Optional[types.ModuleType] = None,
     **kwargs):
@@ -534,7 +522,7 @@ def apply_wrappers(
   """
   if not wrapper_classes:
     wrapper_classes = []
-    for _, c in object_utils.JSONConvertible.registered_types():
+    for _, c in utils.JSONConvertible.registered_types():
       if (issubclass(c, ClassWrapper)
           and c not in (ClassWrapper, _SubclassedWrapperBase)
           and (not where or where(c))
@@ -547,23 +535,30 @@ def _extract_init_signature(
     cls,
     arg_specs=None,
     auto_doc: bool = False,
-    auto_typing: bool = False):
+    auto_typing: bool = False
+) -> Tuple[List[str], List[pg_typing.Field]]:
   """Extract argument fields from class __init__ method."""
   init_method = getattr(cls, '__orig_init__', cls.__init__)
-
-  description = None
-  args_docstr = dict()
+  docstr = None
   if auto_doc:
     # Read args docstr from both class doc string and __init__ doc string.
+    args_docstr = dict()
     if cls.__doc__:
-      cls_docstr = object_utils.DocStr.parse(cls.__doc__)
-      description = schema_utils.schema_description_from_docstr(
-          cls_docstr, include_long_description=True)
+      cls_docstr = utils.DocStr.parse(cls.__doc__)
       args_docstr = cls_docstr.args
     if init_method.__doc__:
-      init_docstr = object_utils.DocStr.parse(init_method.__doc__)
+      init_docstr = utils.DocStr.parse(init_method.__doc__)
       args_docstr.update(init_docstr.args)
-
+    docstr = utils.DocStr(
+        utils.DocStrStyle.GOOGLE,
+        short_description=None,
+        long_description=None,
+        examples=[],
+        args=args_docstr,
+        returns=None,
+        raises=[],
+        blank_after_short_description=True,
+    )
   if init_method is object.__init__:
     if arg_specs:
       raise ValueError(
@@ -572,13 +567,21 @@ def _extract_init_signature(
     init_arg_list = []
     arg_fields = []
   else:
-    signature = pg_typing.get_signature(init_method, auto_typing=auto_typing)
+    signature = pg_typing.Signature.from_signature(
+        inspect.signature(init_method),
+        name=cls.__name__,
+        callable_type=pg_typing.CallableType.METHOD,
+        module_name=cls.__module__,
+        qualname=cls.__qualname__,
+        auto_typing=auto_typing,
+        docstr=docstr,
+    ).annotate(arg_specs)
     if not signature.args or signature.args[0].name != 'self':
       raise ValueError(
           f'{cls.__name__}.__init__ must have `self` as the first argument.')
     # Remove field for 'self'.
-    arg_fields = pg_typing.get_arg_fields(signature, arg_specs, args_docstr)[1:]
+    arg_fields = signature.fields(remove_self=True)
     init_arg_list = [arg.name for arg in signature.args[1:]]
     if signature.varargs is not None:
       init_arg_list.append(f'*{signature.varargs.name}')
-  return (description, init_arg_list, arg_fields)
+  return (init_arg_list, arg_fields)

pyglove/core/symbolic/class_wrapper_test.py

@@ -588,7 +588,7 @@ class ClassWrapperTest(unittest.TestCase):
     self.assertIsInstance(C(1, 2), ClassWrapper)
     self.assertTrue(pg_eq(C(1, 2), C(1, 2)))
     self.assertEqual(list(C.__schema__.fields.keys()), ['x', 'y'])
-    self.assertEqual(repr(C), f'<class {C.type_name!r}>')
+    self.assertEqual(repr(C), f'<class {C.__type_name__!r}>')
 
   def test_custom_metaclass(self):
 
@@ -604,7 +604,7 @@ class ClassWrapperTest(unittest.TestCase):
     A1 = pg_wrap(A)  # pylint: disable=invalid-name
     self.assertTrue(issubclass(A1, ClassWrapper))
     self.assertTrue(issubclass(A1, A))
-    self.assertEqual(A1.type_name, 'pyglove.core.symbolic.class_wrapper_test.A')
+    self.assertEqual(A1.__type_name__, 'pyglove.core.symbolic.class_wrapper_test.A')
     self.assertEqual(A1.__schema__, pg_typing.Schema([]))
     self.assertEqual(A1.foo, 'foo')
     self.assertRegex(repr(A1), r'Symbolic\[.*\]')

pyglove/core/symbolic/compounding.py

@@ -17,10 +17,9 @@ import abc
 import inspect
 import sys
 import types
-from typing import Any, List, Optional, Tuple, Type, Union
+from typing import Any, Dict, List, Optional, Tuple, Type, Union
 
-from pyglove.core import object_utils
-from pyglove.core.symbolic import schema_utils
+from pyglove.core import utils
 from pyglove.core.symbolic.base import Symbolic
 from pyglove.core.symbolic.object import Object
 import pyglove.core.typing as pg_typing
@@ -40,7 +39,7 @@ class Compound(Object):
     # from the user class to compound with.
     Object.__init_subclass__(cls)
 
-  @object_utils.explicit_method_override
+  @utils.explicit_method_override
   def __init__(self, *args, **kwargs):
     # `explicit_init` allows the `__init__` of the other classes that sit after
     # `Compound` to be bypassed.
@@ -53,16 +52,11 @@ _COMPOUND_OWNED_ATTR_NAMES = frozenset(dir(Compound))
 def compound_class(
     factory_fn: types.FunctionType,
     base_class: Optional[Type[Object]] = None,
-    args: Optional[
-        List[
-            Union[
-                Tuple[Tuple[str, pg_typing.KeySpec], pg_typing.ValueSpec, str],
-                Tuple[
-                    Tuple[str, pg_typing.KeySpec], pg_typing.ValueSpec, str, Any
-                ],
-            ]
-        ]
-    ] = None,  # pylint: disable=bad-continuation
+    args: Union[
+        List[Union[pg_typing.Field, pg_typing.FieldDef]],
+        Dict[pg_typing.FieldKeyDef, pg_typing.FieldValueDef],
+        None
+    ] = None,
     *,
     lazy_build: bool = True,
     auto_doc: bool = True,
@@ -127,7 +121,7 @@ def compound_class(
   if not inspect.isfunction(factory_fn):
     raise TypeError('Decorator `compound` is only applicable to functions.')
 
-  schema = schema_utils.function_schema(
+  schema = pg_typing.schema(
       factory_fn,
       args=args,
       returns=pg_typing.Object(base_class) if base_class else None,

pyglove/core/symbolic/compounding_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.compounding."""
-
 import abc
 import dataclasses
 import sys
 import unittest
 
-from pyglove.core import object_utils
 from pyglove.core import typing as pg_typing
+from pyglove.core import utils
 from pyglove.core.symbolic.compounding import compound as pg_compound
 from pyglove.core.symbolic.compounding import compound_class as pg_compound_class
 from pyglove.core.symbolic.dict import Dict
@@ -145,7 +143,7 @@ class UserClassTest(unittest.TestCase):
     class A(Object):
       x: int
 
-      @object_utils.explicit_method_override
+      @utils.explicit_method_override
       def __init__(self, x):
         super().__init__(x=x)
         assert type(self) is A  # pylint: disable=unidiomatic-typecheck