From 328c9566123aa55686c578bc43f83c4369b9615e Mon Sep 17 00:00:00 2001 From: Daiyi Peng Date: Fri, 12 Apr 2024 22:20:59 -0700 Subject: [PATCH] Support `force_dict` flag for `pg.to_json`/`pg.save`. PiperOrigin-RevId: 624396731 --- pyglove/core/object_utils/json_conversion.py | 37 ++++++++++++------- .../core/object_utils/json_conversion_test.py | 4 ++ pyglove/core/symbolic/base.py | 15 +++++--- pyglove/core/symbolic/object_test.py | 11 ++++++ 4 files changed, 48 insertions(+), 19 deletions(-) diff --git a/pyglove/core/object_utils/json_conversion.py b/pyglove/core/object_utils/json_conversion.py index 5d28dfa..5fc4914 100644 --- a/pyglove/core/object_utils/json_conversion.py +++ b/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,6 +348,8 @@ 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. @@ -355,35 +357,42 @@ def to_json(value: Any, **kwargs) -> Any: 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, diff --git a/pyglove/core/object_utils/json_conversion_test.py b/pyglove/core/object_utils/json_conversion_test.py index b3c738c..16b75ee 100644 --- a/pyglove/core/object_utils/json_conversion_test.py +++ b/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__', diff --git a/pyglove/core/symbolic/base.py b/pyglove/core/symbolic/base.py index ca25d74..88f8762 100644 --- a/pyglove/core/symbolic/base.py +++ b/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, diff --git a/pyglove/core/symbolic/object_test.py b/pyglove/core/symbolic/object_test.py index b099fa9..4248692 100644 --- a/pyglove/core/symbolic/object_test.py +++ b/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), {