Add basic support for unions using sdfChoice

onedm/sdf/from_type.py

@@ -1,6 +1,6 @@
 """Conversion from native types to sdfData."""
 
-from typing import Type
+from typing import Any, Type
 
 from pydantic import TypeAdapter
 from pydantic.json_schema import GenerateJsonSchema
@@ -59,11 +59,12 @@ def enum_schema(self, schema: core_schema.EnumSchema):
             del definition["enum"]
         return definition
 
-    def literal_schema(self, schema: core_schema.LiteralSchema):
-        definition = super().literal_schema(schema)
-        if "const" in definition and "enum" in definition:
-            # Don't need both
-            del definition["enum"]
+    def union_schema(self, schema: core_schema.UnionSchema):
+        definition = super().union_schema(schema)
+        definition["sdfChoice"] = {
+            f"choice-{i}": choice for i, choice in enumerate(definition["anyOf"], start=1)
+        }
+        del definition["anyOf"]
         return definition
 
     def bytes_schema(self, schema: core_schema.BytesSchema):

tests/sdf/test_from_type.py

@@ -99,6 +99,16 @@ class MyEnum(str, enum.Enum):
     assert not data.nullable
 
 
+def test_union():
+    data = data_from_type(int | str)
+
+    assert len(data.choices) == 2
+    assert "choice-1" in data.choices
+    assert "choice-2" in data.choices
+    assert data.choices["choice-1"].type == "integer"
+    assert data.choices["choice-2"].type == "string"
+
+
 def test_const():
     data = data_from_type(Literal["const"])