Honor __serialization_key__ when serializing schema to JSON.

PiperOrigin-RevId: 572067682
google · Oct 9, 2023 · cff81b6 · cff81b6
1 parent 1ed40e7
commit cff81b6
Show file tree

Hide file tree

Showing 2 changed files with 7 additions and 6 deletions.
diff --git a/langfun/core/structured/schema.py b/langfun/core/structured/schema.py
@@ -150,7 +150,7 @@ def _node(vs: pg.typing.ValueSpec) -> Any:
         return [_node(vs.element.value)]
       elif isinstance(vs, pg.typing.Object):
         if issubclass(vs.cls, pg.Object):
-          d = {pg.JSONConvertible.TYPE_NAME_KEY: vs.cls.__type_name__}
+          d = {pg.JSONConvertible.TYPE_NAME_KEY: vs.cls.__serialization_key__}
           d.update(
               {
                   str(k): _node(f.value)

diff --git a/langfun/core/structured/schema_test.py b/langfun/core/structured/schema_test.py
@@ -35,6 +35,9 @@ class Itinerary(pg.Object):
   ]
 
 
+Itinerary.__serialization_key__ = 'Itinerary'
+
+
 class Node(pg.Object):
   children: list['Node']
 
@@ -100,7 +103,7 @@ def test_schema_dict(self):
             'result': [
                 {
                     'x': {
-                        '_type': Itinerary.__type_name__,
+                        '_type': 'Itinerary',
                         'day': pg.typing.Int(min_value=1),
                         'activities': [{
                             '_type': Activity.__type_name__,
@@ -139,11 +142,10 @@ def test_schema_repr(self):
     self.assertEqual(
         schema.schema_str(protocol='json'),
         (
-            '{"result": [{"x": {"_type": "%s", "day":'
+            '{"result": [{"x": {"_type": "Itinerary", "day":'
             ' int(min=1), "type": "daytime" | "nighttime", "activities":'
             ' [{"_type": "%s", "description": str}], "hotel":'
             ' str(regex=.*Hotel) | None}}]}' % (
-                Itinerary.__type_name__,
                 Activity.__type_name__,
             )
         ),
@@ -495,11 +497,10 @@ def test_repr(self):
     self.assertEqual(
         schema_lib.SchemaJsonRepr().repr(schema),
         (
-            '{"result": [{"x": {"_type": "%s", "day":'
+            '{"result": [{"x": {"_type": "Itinerary", "day":'
             ' int(min=1), "type": "daytime" | "nighttime", "activities":'
             ' [{"_type": "%s", "description": str}], "hotel":'
             ' str(regex=.*Hotel) | None}}]}' % (
-                Itinerary.__type_name__,
                 Activity.__type_name__,
             )
         ),