Skip to content

Commit 787b7f4

Browse files
fix (serialization) : encode unknown dataclasses as plain dicts;add unit test
Signed-off-by: Gautam Datla <[email protected]>
1 parent 89225dc commit 787b7f4

File tree

2 files changed

+49
-6
lines changed

2 files changed

+49
-6
lines changed

nemoguardrails/colang/v2_x/runtime/serialization.py

Lines changed: 14 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -86,13 +86,21 @@ def encode_to_dict(obj: Any, refs: Dict[int, Any]):
8686
"value": {k: encode_to_dict(v, refs) for k, v in obj.items()},
8787
}
8888
elif is_dataclass(obj):
89-
value = {
90-
"__type": type(obj).__name__,
91-
"value": {
92-
k: encode_to_dict(getattr(obj, k), refs)
93-
for k in obj.__dataclass_fields__.keys()
94-
},
89+
# Encode dataclasses. If it's a known class (present in name_to_class),
90+
# keep its type tag so we can fully round-trip via json_to_state.
91+
# Otherwise, fall back to a plain dict to avoid "Unknown d_type" on decode.
92+
cls = type(obj)
93+
encoded_fields = {
94+
k: encode_to_dict(getattr(obj, k), refs)
95+
for k in obj.__dataclass_fields__.keys()
9596
}
97+
98+
if cls.__name__ in name_to_class and name_to_class[cls.__name__] is cls:
99+
value = {"__type": cls.__name__, "value": encoded_fields}
100+
else:
101+
# Unknown dataclass → JSON-friendly dict
102+
value = {"__type": "dict", "value": encoded_fields}
103+
96104
elif isinstance(obj, RailsConfig):
97105
value = {
98106
"__type": "RailsConfig",
Lines changed: 35 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,35 @@
1+
# SPDX-FileCopyrightText: Copyright (c) 2023 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
2+
# SPDX-License-Identifier: Apache-2.0
3+
#
4+
# Licensed under the Apache License, Version 2.0 (the "License");
5+
# you may not use this file except in compliance with the License.
6+
# You may obtain a copy of the License at
7+
#
8+
# http://www.apache.org/licenses/LICENSE-2.0
9+
#
10+
# Unless required by applicable law or agreed to in writing, software
11+
# distributed under the License is distributed on an "AS IS" BASIS,
12+
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13+
# See the License for the specific language governing permissions and
14+
# limitations under the License.
15+
16+
import json
17+
from dataclasses import dataclass
18+
19+
from nemoguardrails.colang.v2_x.runtime.serialization import state_to_json
20+
21+
22+
@dataclass
23+
class Foo:
24+
bar: str
25+
baz: int
26+
27+
28+
def test_state_to_json_unknown_dataclass_encodes_as_dict():
29+
js = state_to_json({"out": Foo("ok", 1)})
30+
d = json.loads(js)
31+
# We expect unknown dataclasses to be encoded as plain dicts
32+
assert d["__type"] == "dict"
33+
out = d["value"]["out"]
34+
assert out["__type"] == "dict"
35+
assert out["value"] == {"bar": "ok", "baz": 1}

0 commit comments

Comments
 (0)