pg.to_json to support opaque object serialization with pickle.

PiperOrigin-RevId: 559962280

Author: daiyip

pyglove/core/object_utils/json_conversion.py

@@ -19,6 +19,7 @@
 import importlib
 import inspect
 import marshal
+import pickle
 import types
 import typing
 from typing import Any, Callable, Dict, Iterable, List, Optional, Tuple, Type, TypeVar, Union
@@ -199,8 +200,8 @@ def register(
     Args:
       type_name: A global unique string identifier for subclass.
       subclass: A subclass of JSONConvertible.
-      override_existing: If True, registering an type name is allowed.
-        Otherwise an error will be raised.
+      override_existing: If True, override the class if the type name is
+        already present in the registry. Otherwise an error will be raised.
     """
     cls._TYPE_REGISTRY.register(type_name, subclass, override_existing)
 
@@ -246,6 +247,51 @@ def __init_subclass__(cls):
       JSONConvertible.register(type_name, cls, override_existing=True)
 
 
+def _type_name(type_or_function: Union[Type[Any], types.FunctionType]) -> str:
+  """Returns the ID for a type or function."""
+  return f'{type_or_function.__module__}.{type_or_function.__qualname__}'
+
+
+class _OpaqueObject(JSONConvertible):
+  """An JSON converter for opaque Python objects."""
+
+  def __init__(self, value: Any, encoded: bool = False):
+    if encoded:
+      value = self.decode(value)
+    self._value = value
+
+  @property
+  def value(self) -> Any:
+    """Returns the decoded value."""
+    return self._value
+
+  def encode(self, value: Any) -> JSONValueType:
+    try:
+      return base64.encodebytes(pickle.dumps(value)).decode('utf-8')
+    except Exception as e:
+      raise ValueError(
+          f'Cannot encode opaque object {value!r} with pickle.') from e
+
+  def decode(self, json_value: JSONValueType) -> Any:
+    assert isinstance(json_value, str), json_value
+    try:
+      return pickle.loads(base64.decodebytes(json_value.encode('utf-8')))
+    except Exception as e:
+      raise ValueError('Cannot decode opaque object with pickle.') from e
+
+  def to_json(self, **kwargs) -> JSONValueType:
+    return self.to_json_dict({
+        'value': self.encode(self._value)
+    }, **kwargs)
+
+  @classmethod
+  def from_json(cls, json_value: JSONValueType, *args, **kwargs) -> Any:
+    del args, kwargs
+    assert isinstance(json_value, dict) and 'value' in json_value, json_value
+    encoder = cls(json_value['value'], encoded=True)
+    return encoder.value
+
+
 def registered_types() -> Iterable[Tuple[str, Type[JSONConvertible]]]:
   """Returns an iterator of registered (serialization key, class) tuples."""
   return JSONConvertible.registered_types()
@@ -298,7 +344,7 @@ def to_json(value: Any, **kwargs) -> Any:
       converter = JSONConvertible.TYPE_CONVERTER(type(value))   # pylint: disable=not-callable
       if converter:
         return to_json(converter(value))
-    raise ValueError(f'Cannot convert complex type {value} to JSON.')
+    return _OpaqueObject(value).to_json(**kwargs)
 
 
 def from_json(json_value: JSONValueType) -> Any:
@@ -329,8 +375,6 @@ def from_json(json_value: JSONValueType) -> Any:
       return _function_from_json(json_value)
     elif type_name == 'method':
       return _method_from_json(json_value)
-    # elif type_name == 'annotation':
-    #   return _annotation_from_json(json_value)
     else:
       cls = JSONConvertible.class_from_typename(type_name)
       if cls is None:
@@ -346,11 +390,6 @@ def from_json(json_value: JSONValueType) -> Any:
 #
 
 
-def _type_name(type_or_function: Union[Type[Any], types.FunctionType]) -> str:
-  """Returns the ID for a type or function."""
-  return f'{type_or_function.__module__}.{type_or_function.__qualname__}'
-
-
 def _type_to_json(t: Type[Any]) -> Dict[str, str]:
   """Converts a type to a JSON dict."""
   type_name = _type_name(t)

pyglove/core/object_utils/json_conversion_test.py

@@ -14,6 +14,7 @@
 """Tests for pyglove.object_utils.json_conversion."""
 
 import abc
+import datetime
 import typing
 import unittest
 from pyglove.core.object_utils import json_conversion
@@ -38,6 +39,15 @@ def static_method():
       def instance_method(self):
         return str(self)
 
+  def __init__(self, x):
+    self.x = x
+
+  def __eq__(self, other):
+    return isinstance(other, X) and self.x == other.x
+
+  def __ne__(self, other):
+    return not self.__eq__(other)
+
 
 def bar():
   pass
@@ -256,6 +266,23 @@ def test_json_conversion_for_methods(self):
         ValueError, 'Cannot convert instance method .* to JSON.'):
       json_conversion.to_json(X.Y.Z().instance_method)
 
+  def test_json_conversion_for_opaque_objects(self):
+    self.assert_conversion_equal(X(1))
+    self.assert_conversion_equal(datetime.datetime.now())
+
+    class LocalX:
+      pass
+
+    with self.assertRaisesRegex(
+        ValueError, 'Cannot encode opaque object .* with pickle.'):
+      json_conversion.to_json(LocalX())
+
+    json_dict = json_conversion.to_json(X(1))
+    json_dict['value'] = 'abc'
+    with self.assertRaisesRegex(
+        ValueError, 'Cannot decode opaque object with pickle.'):
+      json_conversion.from_json(json_dict)
+
 
 if __name__ == '__main__':
   unittest.main()

pyglove/core/symbolic/dict_test.py

@@ -1854,7 +1854,7 @@ def __eq__(self, other):
     ])
     sd = Dict(x=1, y=A(2.0), value_spec=spec)
     with self.assertRaisesRegex(
-        ValueError, 'Cannot convert complex type .* to JSON.'):
+        ValueError, 'Cannot encode opaque object .* with pickle'):
       sd.to_json_str()
 
     pg_typing.register_converter(A, float, convert_fn=lambda x: x.value)
@@ -1905,7 +1905,7 @@ class A:
       pass
 
     with self.assertRaisesRegex(
-        ValueError, 'Cannot convert complex type .* to JSON.'):
+        ValueError, 'Cannot encode opaque object .* with pickle'):
       base.to_json(Dict(x=A()))
 
 

pyglove/core/symbolic/list_test.py

@@ -1593,7 +1593,7 @@ def __eq__(self, other):
     ]))
     sl = List([dict(x=1, y=A(2.0))], value_spec=spec)
     with self.assertRaisesRegex(
-        ValueError, 'Cannot convert complex type .* to JSON.'):
+        ValueError, 'Cannot encode opaque object .* with pickle.'):
       sl.to_json_str()
 
     pg_typing.register_converter(A, float, convert_fn=lambda x: x.value)
@@ -1646,7 +1646,7 @@ class A:
       pass
 
     with self.assertRaisesRegex(
-        ValueError, 'Cannot convert complex type .* to JSON.'):
+        ValueError, 'Cannot encode opaque object .* with pickle.'):
       base.to_json(List([A()]))
 
 

pyglove/core/symbolic/object.py

@@ -456,7 +456,9 @@ class Foo(pg.Object):
     Returns:
       A symbolic Object instance.
     """
-    return cls(allow_partial=allow_partial, root_path=root_path, **json_value)
+    return cls(allow_partial=allow_partial, root_path=root_path, **{
+        k: base.from_json(v) for k, v in json_value.items()
+    })
 
   @object_utils.explicit_method_override
   def __init__(

pyglove/core/symbolic/object_test.py

@@ -2874,7 +2874,7 @@ def test_serialization_with_converter(self):
 
     c = self._C(self._X(1), y=True)
     with self.assertRaisesRegex(
-        ValueError, 'Cannot convert complex type .* to JSON.'):
+        ValueError, 'Cannot encode opaque object .* with pickle'):
       c.to_json_str()
 
     pg_typing.register_converter(self._X, int, convert_fn=lambda x: x.value)
@@ -2902,7 +2902,7 @@ class Z:
       pass
 
     with self.assertRaisesRegex(
-        ValueError, 'Cannot convert complex type .* to JSON.'):
+        ValueError, 'Cannot encode opaque object .* with pickle'):
       base.to_json(self._A(w=Z(), y=True))
 
     with self.assertRaisesRegex(