Add optional field for Unions to track types at runtime (#15)

Author: giograno

src/py_avro_schema/_schemas.py

@@ -72,6 +72,8 @@
 
 NamesType = List[str]
 
+RUNTIME_TYPE_KEY = "_runtime_type"
+
 
 class TypeNotSupportedError(TypeError):
     """Error raised when a Avro schema cannot be generated for a given Python type"""
@@ -137,6 +139,9 @@ class Option(enum.Flag):
     # the two cases.
     MARK_NON_TOTAL_TYPED_DICTS = enum.auto()
 
+    #: Adds a _runtime_type field to the record schemas that contains the name of the class
+    ADD_RUNTIME_TYPE_FIELD = enum.auto()
+
 
 JSON_OPTIONS = [opt for opt in Option if opt.name and opt.name.startswith("JSON_")]
 
@@ -1105,6 +1110,13 @@ def _record_field(self, py_field: dataclasses.Field) -> RecordField:
 
         return field_obj
 
+    def data_before_deduplication(self, names: NamesType) -> JSONObj:
+        """Return the schema data"""
+        data = super().data_before_deduplication(names)
+        if Option.ADD_RUNTIME_TYPE_FIELD in self.options:
+            data["fields"].append({"name": RUNTIME_TYPE_KEY, "type": ["null", "string"]})
+        return data
+
 
 @register_schema
 class PydanticSchema(RecordSchema):
@@ -1239,6 +1251,13 @@ def _record_field(self, py_field: tuple[str, Type]) -> RecordField:
         )
         return field_obj
 
+    def data_before_deduplication(self, names: NamesType) -> JSONObj:
+        """Return the schema data"""
+        data = super().data_before_deduplication(names)
+        if Option.ADD_RUNTIME_TYPE_FIELD in self.options:
+            data["fields"].append({"name": RUNTIME_TYPE_KEY, "type": ["null", "string"]})
+        return data
+
 
 @register_schema
 class TypedDictSchema(RecordSchema):

tests/test_avro_schema.py

@@ -18,6 +18,7 @@
 
 import py_avro_schema as pas
 from py_avro_schema._alias import register_type_alias, register_type_aliases
+from py_avro_schema._testing import assert_schema
 
 
 def test_package_has_version():
@@ -64,3 +65,18 @@ class PyTypedDict(TypedDict):
         "test_avro_schema.SuperOldDict",
         "test_avro_schema.VeryOldDict",
     ]
+
+
+def test_add_type_field():
+    class PyType:
+        field: str
+
+    expected = {
+        "type": "record",
+        "name": "PyType",
+        "fields": [
+            {"name": "field", "type": "string"},
+            {"name": "_runtime_type", "type": ["null", "string"]},
+        ],
+    }
+    assert_schema(PyType, expected, options=pas.Option.ADD_RUNTIME_TYPE_FIELD)