Decouple JSON conversion and type conversion.

- Removes the concept of json value converters, as it is no longer needed because the introduction of opaque object serialization.
- Accepts tuple as the value for the `dest` argument for `pg.typing.get_converter`, and remove `pg.typing.get_first_applicable_converter`.

PiperOrigin-RevId: 560116998

Author: daiyip

pyglove/core/object_utils/json_conversion.py

@@ -22,7 +22,7 @@
 import pickle
 import types
 import typing
-from typing import Any, Callable, Dict, Iterable, List, Optional, Tuple, Type, TypeVar, Union
+from typing import Any, Dict, Iterable, List, Optional, Tuple, Type, TypeVar, Union
 
 # Nestable[T] is a (maybe) nested structure of T, which could be T, a Dict
 # a List or a Tuple of Nestable[T]. We use a Union to fool PyType checker to
@@ -139,12 +139,6 @@ def from_json(cls, json_value, **kwargs):
   # Marker (as the first element of a list) for serializing tuples.
   TUPLE_MARKER = '__tuple__'
 
-  # Type converter that converts a complex type to basic JSON value type.
-  # When this field is set by users, the converter will be invoked when a
-  # complex value cannot be serialized by existing methods.
-  TYPE_CONVERTER: Optional[
-      Callable[[Type[Any]], Callable[[Any], JSONValueType]]] = None
-
   # Class property that indicates whether to automatically register class
   # for deserialization. Subclass can override.
   auto_register = True
@@ -340,10 +334,6 @@ def to_json(value: Any, **kwargs) -> Any:
   elif value is ...:
     return {JSONConvertible.TYPE_NAME_KEY: 'type', 'name': 'builtins.Ellipsis'}
   else:
-    if JSONConvertible.TYPE_CONVERTER is not None:
-      converter = JSONConvertible.TYPE_CONVERTER(type(value))   # pylint: disable=not-callable
-      if converter:
-        return to_json(converter(value))
     return _OpaqueObject(value).to_json(**kwargs)
 
 

pyglove/core/symbolic/dict_test.py

@@ -1838,34 +1838,6 @@ def test_schematized(self):
         ]))
     self.assertEqual(sd.to_json_str(), '{"x": 1, "y": null}')
 
-  def test_serialization_with_converter(self):
-
-    class A:
-
-      def __init__(self, value: float):
-        self.value = value
-
-      def __eq__(self, other):
-        return isinstance(other, A) and other.value == self.value
-
-    spec = pg_typing.Dict([
-        ('x', pg_typing.Int()),
-        ('y', pg_typing.Object(A))
-    ])
-    sd = Dict(x=1, y=A(2.0), value_spec=spec)
-    with self.assertRaisesRegex(
-        ValueError, 'Cannot encode opaque object .* with pickle'):
-      sd.to_json_str()
-
-    pg_typing.register_converter(A, float, convert_fn=lambda x: x.value)
-    pg_typing.register_converter(float, A, convert_fn=A)
-
-    self.assertEqual(sd.to_json(), {'x': 1, 'y': 2.0})
-    self.assertEqual(Dict.from_json(sd.to_json(), value_spec=spec), sd)
-
-    self.assertEqual(sd.to_json_str(), '{"x": 1, "y": 2.0}')
-    self.assertEqual(base.from_json_str(sd.to_json_str(), value_spec=spec), sd)
-
   def test_hide_default_values(self):
     sd = Dict.partial(
         x=1,

pyglove/core/symbolic/list_test.py

@@ -1577,34 +1577,6 @@ def test_schematized(self):
         ])))
     self.assertEqual(sl.to_json_str(), '[{"x": 1, "y": null}]')
 
-  def test_serialization_with_converter(self):
-
-    class A:
-
-      def __init__(self, value: float):
-        self.value = value
-
-      def __eq__(self, other):
-        return isinstance(other, A) and other.value == self.value
-
-    spec = pg_typing.List(pg_typing.Dict([
-        ('x', pg_typing.Int()),
-        ('y', pg_typing.Object(A))
-    ]))
-    sl = List([dict(x=1, y=A(2.0))], value_spec=spec)
-    with self.assertRaisesRegex(
-        ValueError, 'Cannot encode opaque object .* with pickle.'):
-      sl.to_json_str()
-
-    pg_typing.register_converter(A, float, convert_fn=lambda x: x.value)
-    pg_typing.register_converter(float, A, convert_fn=A)
-
-    self.assertEqual(sl.to_json(), [{'x': 1, 'y': 2.0}])
-    self.assertEqual(List.from_json(sl.to_json(), value_spec=spec), sl)
-
-    self.assertEqual(sl.to_json_str(), '[{"x": 1, "y": 2.0}]')
-    self.assertEqual(base.from_json_str(sl.to_json_str(), value_spec=spec), sl)
-
   def test_serialization_with_tuple(self):
     self.assertEqual(base.to_json_str((1, 2)), '["__tuple__", 1, 2]')
     self.assertEqual(base.from_json_str('["__tuple__", 1]'), (1,))

pyglove/core/symbolic/object_test.py

@@ -2870,21 +2870,6 @@ def from_json(cls, json_dict, *args, **kwargs):
     a = self._A(Y(1), y=True)
     self.assertEqual(base.from_json_str(a.to_json_str()), a)
 
-  def test_serialization_with_converter(self):
-
-    c = self._C(self._X(1), y=True)
-    with self.assertRaisesRegex(
-        ValueError, 'Cannot encode opaque object .* with pickle'):
-      c.to_json_str()
-
-    pg_typing.register_converter(self._X, int, convert_fn=lambda x: x.value)
-    pg_typing.register_converter(int, self._X, convert_fn=self._X)
-
-    self.assertEqual(
-        c.to_json_str(),
-        '{"_type": "%s", "w": 1, "x": 1, "y": true}' % self._C.type_name)
-    self.assertEqual(base.from_json_str(c.to_json_str()), c)
-
   def test_hide_default_values(self):
     b = self._B('foo', 1)
     self.assertEqual(

pyglove/core/typing/__init__.py

@@ -356,8 +356,6 @@ class Foo(pg.Object):
 # Type conversion.
 from pyglove.core.typing.type_conversion import register_converter
 from pyglove.core.typing.type_conversion import get_converter
-from pyglove.core.typing.type_conversion import get_first_applicable_converter
-from pyglove.core.typing.type_conversion import get_json_value_converter
 
 # Inspect helpers.
 from pyglove.core.typing.inspect import is_subclass

pyglove/core/typing/type_conversion.py

@@ -27,8 +27,6 @@ class _TypeConverterRegistry:
   def __init__(self):
     """Constructor."""
     self._converter_list = []
-    self._json_value_types = set(
-        [int, float, bool, type(None), list, tuple, dict, str])
 
   def register(
       self,
@@ -44,11 +42,7 @@ def register(
     ):
       raise TypeError('Argument \'src\' and \'dest\' must be a type or '
                       'tuple of types.')
-    if isinstance(dest, tuple):
-      json_value_convertible = any(d in self._json_value_types for d in dest)
-    else:
-      json_value_convertible = dest in self._json_value_types
-    self._converter_list.append((src, dest, convert_fn, json_value_convertible))
+    self._converter_list.append((src, dest, convert_fn))
 
   def get_converter(
       self, src: Type[Any], dest: Type[Any]) -> Optional[Callable[[Any], Any]]:
@@ -58,51 +52,30 @@ def get_converter(
     # We may consider more efficient way to do lookup in future.
     # NOTE(daiyip): We do reverse lookup since usually subclass converter
     # is register after base class.
-    for src_type, dest_type, converter, _ in reversed(self._converter_list):
+    for src_type, dest_type, converter in reversed(self._converter_list):
       if pg_inspect.is_subclass(src, src_type):
         dest_types = dest_type if isinstance(dest_type, tuple) else (dest_type,)
         for dest_type in dest_types:
           if pg_inspect.is_subclass(dest_type, dest):
             return converter
     return None
 
-  def get_json_value_converter(
-      self, src: Type[Any]) -> Optional[Callable[[Any], Any]]:
-    """Get converter from source type to a JSON simple type."""
-    for src_type, _, converter, json_value_convertible in reversed(
-        self._converter_list):
-      if pg_inspect.is_subclass(src, src_type) and json_value_convertible:
-        return converter
-    return None
-
 
 _TYPE_CONVERTER_REGISTRY = _TypeConverterRegistry()
 
 
 def get_converter(
-    src: Type[Any], dest: Type[Any]
+    src: Type[Any], dest: Union[Type[Any], Tuple[Type[Any], ...]]
 ) -> Optional[Callable[[Any], Any]]:
   """Get converter from source type to destination type."""
-  return _TYPE_CONVERTER_REGISTRY.get_converter(src, dest)
-
-
-def get_first_applicable_converter(
-    src_type: Type[Any],
-    dest_type: Union[Type[Any], Tuple[Type[Any], ...]]):
-  """Get first applicable converter."""
-  dest_types = dest_type if isinstance(dest_type, tuple) else (dest_type,)
-  for dest_type in dest_types:
-    converter = get_converter(src_type, dest_type)
+  dest_types = dest if isinstance(dest, tuple) else (dest,)
+  for dest in dest_types:
+    converter = _TYPE_CONVERTER_REGISTRY.get_converter(src, dest)
     if converter is not None:
       return converter
   return None
 
 
-def get_json_value_converter(src: Type[Any]) -> Optional[Callable[[Any], Any]]:
-  """Get converter from source type to a JSON simple type."""
-  return _TYPE_CONVERTER_REGISTRY.get_json_value_converter(src)
-
-
 def register_converter(
     src_type: Union[Type[Any], Tuple[Type[Any], ...]],
     dest_type: Union[Type[Any], Tuple[Type[Any], ...]],
@@ -148,4 +121,3 @@ def _register_builtin_converters():
 
 
 _register_builtin_converters()
-object_utils.JSONConvertible.TYPE_CONVERTER = get_json_value_converter

pyglove/core/typing/type_conversion_test.py

@@ -45,18 +45,13 @@ def __init__(self, x, y):
 
     self.assertIs(type_conversion.get_converter(A, str), a_converter)
     self.assertIs(type_conversion.get_converter(A, int), a_converter)
-    self.assertIs(type_conversion.get_json_value_converter(A), a_converter)
 
-    self.assertIsNone(
-        type_conversion.get_first_applicable_converter(A, (float, bool)))
-    self.assertIs(
-        type_conversion.get_first_applicable_converter(A, (float, int)),
-        a_converter)
+    self.assertIsNone(type_conversion.get_converter(A, (float, bool)))
+    self.assertIs(type_conversion.get_converter(A, (float, int)), a_converter)
 
     # B is a subclass of A, so A's converter applies.
     self.assertIs(type_conversion.get_converter(B, str), a_converter)
     self.assertIs(type_conversion.get_converter(B, int), a_converter)
-    self.assertIs(type_conversion.get_json_value_converter(B), a_converter)
 
     b_converter = lambda b: b.y
     type_conversion.register_converter(B, (str, int), b_converter)
@@ -65,13 +60,9 @@ def __init__(self, x, y):
     # match.
     self.assertIs(type_conversion.get_converter(B, str), b_converter)
     self.assertIs(type_conversion.get_converter(B, int), b_converter)
-    self.assertIs(type_conversion.get_json_value_converter(B), b_converter)
 
-    self.assertIsNone(
-        type_conversion.get_first_applicable_converter(B, (float, bool)))
-    self.assertIs(
-        type_conversion.get_first_applicable_converter(B, (float, int)),
-        b_converter)
+    self.assertIsNone(type_conversion.get_converter(B, (float, bool)))
+    self.assertIs(type_conversion.get_converter(B, (float, int)), b_converter)
 
     # Test generics.
     T = typing.TypeVar('T')
@@ -120,8 +111,6 @@ class B:
         vs.Union([vs.Object(B), vs.Float()]).apply(A(1)).x, 1)
     self.assertEqual(vs.Object(A).apply(1).x, 1)
     self.assertEqual(vs.Object(A).apply('foo').x, 'foo')
-    self.assertEqual(type_conversion.get_json_value_converter(A)(A(1)), 1)
-    self.assertIsNone(type_conversion.get_json_value_converter(B))
 
 
 class BuiltInConversionsTest(unittest.TestCase):
@@ -136,9 +125,6 @@ def test_datetime_to_int(self):
     now = datetime.datetime.utcfromtimestamp(timestamp)
     self.assertEqual(vs.Object(datetime.datetime).apply(timestamp), now)
     self.assertEqual(vs.Int().apply(now), timestamp)
-    self.assertEqual(
-        type_conversion.get_json_value_converter(datetime.datetime)(now),
-        timestamp)
 
   def test_keypath_to_str(self):
     """Test built-in converter between string and KeyPath."""
@@ -151,9 +137,6 @@ def test_keypath_to_str(self):
         ['a', 'b', 'c'])
     self.assertEqual(
         vs.Str().apply(object_utils.KeyPath.parse('a.b.c')), 'a.b.c')
-    self.assertEqual(
-        type_conversion.get_json_value_converter(object_utils.KeyPath)(
-            object_utils.KeyPath.parse('a.b.c')), 'a.b.c')
 
 
 if __name__ == '__main__':

pyglove/core/typing/value_specs.py

@@ -261,8 +261,7 @@ def apply(
         and self.value_type is not None
         and not pg_inspect.is_instance(value, self.value_type)
     ):
-      converter = type_conversion.get_first_applicable_converter(
-          type(value), self.value_type)
+      converter = type_conversion.get_converter(type(value), self.value_type)
       if converter is None:
         raise TypeError(
             object_utils.message_on_path(
@@ -2500,8 +2499,11 @@ def _apply(self,
     # tf.Tensor should be able to accept tf.Variable.
     matched_candidate = None
     for c in self._candidates:
-      if c.type_resolved and type_conversion.get_first_applicable_converter(
-          type(value), c.value_type) is not None:
+      if (
+          c.type_resolved
+          and type_conversion.get_converter(type(value), c.value_type)
+          is not None
+      ):
         matched_candidate = c
         break