Support force_dict flag for pg.to_json/pg.save.

PiperOrigin-RevId: 624361438

Author: daiyip

pyglove/core/object_utils/json_conversion.py

@@ -336,7 +336,7 @@ def registered_types() -> Iterable[Tuple[str, Type[JSONConvertible]]]:
   return JSONConvertible.registered_types()
 
 
-def to_json(value: Any, **kwargs) -> Any:
+def to_json(value: Any, *, force_dict: bool = False, **kwargs) -> Any:
   """Serializes a (maybe) JSONConvertible value into a plain Python object.
 
   Args:
@@ -348,42 +348,51 @@ def to_json(value: Any, **kwargs) -> Any:
       * Tuple types;
       * Dict types.
 
+    force_dict: If True, "_type" keys will be renamed to "type_name".
+      As a result, JSONConvertible objects will be returned as dict.
     **kwargs: Keyword arguments to pass to value.to_json if value is
       JSONConvertible.
 
   Returns:
     JSON value.
   """
   if isinstance(value, (type(None), bool, int, float, str)):
-    return value
+    v = value
   elif isinstance(value, JSONConvertible):
-    return value.to_json(**kwargs)
+    v = value.to_json(**kwargs)
   elif isinstance(value, tuple):
-    return [JSONConvertible.TUPLE_MARKER] + to_json(list(value), **kwargs)
+    v = [JSONConvertible.TUPLE_MARKER] + to_json(list(value), **kwargs)
   elif isinstance(value, list):
-    return [to_json(item, **kwargs) for item in value]
+    v = [to_json(item, **kwargs) for item in value]
   elif isinstance(value, dict):
-    return {k: to_json(v, **kwargs) for k, v in value.items()}
+    v = {k: to_json(v, **kwargs) for k, v in value.items()}
   elif isinstance(value, (type, typing.GenericAlias)):  # pytype: disable=module-attr
-    return _type_to_json(value)
+    v = _type_to_json(value)
   elif inspect.isbuiltin(value):
-    return _builtin_function_to_json(value)
+    v = _builtin_function_to_json(value)
   elif inspect.isfunction(value):
-    return _function_to_json(value)
+    v = _function_to_json(value)
   elif inspect.ismethod(value):
-    return _method_to_json(value)
+    v = _method_to_json(value)
   # pytype: disable=module-attr
   elif isinstance(value, typing._Final):  # pylint: disable=protected-access
     # pytype: enable=module-attr
-    return _annotation_to_json(value)
+    v = _annotation_to_json(value)
   elif value is ...:
-    return {JSONConvertible.TYPE_NAME_KEY: 'type', 'name': 'builtins.Ellipsis'}
+    v = {JSONConvertible.TYPE_NAME_KEY: 'type', 'name': 'builtins.Ellipsis'}
   else:
+    v, converted = None, False
     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)
+        v = to_json(converter(value))
+        converted = True
+    if not converted:
+      v = _OpaqueObject(value).to_json(**kwargs)
+
+  if force_dict:
+    v = replace_type_with_type_names(v)
+  return v
 
 
 def from_json(json_value: JSONValueType,

pyglove/core/object_utils/json_conversion_test.py

@@ -298,6 +298,10 @@ class LocalX:
       json_conversion.from_json(json_dict)
 
   def test_json_conversion_force_dict(self):
+    self.assertEqual(
+        json_conversion.to_json([int], force_dict=True),
+        [{'name': 'builtins.int', 'type_name': 'type'}]
+    )
     self.assertEqual(
         json_conversion.from_json([
             '__tuple__',

pyglove/core/symbolic/base.py

@@ -952,11 +952,11 @@ def clone(
 
   def to_json(self, **kwargs) -> object_utils.JSONValueType:
     """Alias for `sym_jsonify`."""
-    return self.sym_jsonify(**kwargs)
+    return to_json(self, **kwargs)
 
   def to_json_str(self, json_indent: Optional[int] = None, **kwargs) -> str:
     """Serializes current object into a JSON string."""
-    return json.dumps(self.sym_jsonify(**kwargs), indent=json_indent)
+    return to_json_str(self, json_indent=json_indent, **kwargs)
 
   @classmethod
   def load(cls, *args, **kwargs) -> Any:
@@ -2114,7 +2114,7 @@ class A(pg.Object):
       **kwargs)
 
 
-def to_json(value: Any, **kwargs) -> Any:
+def to_json(value: Any, *, force_dict: bool = False, **kwargs) -> Any:
   """Serializes a (maybe) symbolic value into a plain Python object.
 
   Example::
@@ -2139,6 +2139,8 @@ class A(pg.Object):
       * Tuple types;
       * Dict types.
 
+    force_dict: If True, "_type" keys will be renamed to "type_name".
+      As a result, JSONConvertible objects will be returned as dict.
     **kwargs: Keyword arguments to pass to value.to_json if value is
       JSONConvertible.
 
@@ -2148,8 +2150,11 @@ class A(pg.Object):
   # NOTE(daiyip): special handling `sym_jsonify` since symbolized
   # classes may have conflicting `to_json` method in their existing classes.
   if isinstance(value, Symbolic):
-    return value.sym_jsonify(**kwargs)
-  return object_utils.to_json(value, **kwargs)
+    v = value.sym_jsonify(**kwargs)
+    if force_dict:
+      v = object_utils.json_conversion.replace_type_with_type_names(v)
+    return v
+  return object_utils.to_json(value, force_dict=force_dict, **kwargs)
 
 
 def to_json_str(value: Any,

pyglove/core/symbolic/object_test.py

@@ -2944,6 +2944,17 @@ class Q(Object):
       p: P
       y: str
 
+    self.assertEqual(
+        Q(P(1), y='foo').to_json(force_dict=True),
+        {
+            'p': {
+                'type_name': P.__type_name__,
+                'x': 1
+            },
+            'y': 'foo',
+            'type_name': Q.__type_name__,
+        }
+    )
     self.assertEqual(
         base.from_json_str(Q(P(1), y='foo').to_json_str(), force_dict=True),
         {