Fix Python schema conversion for variadic tuples with union elements

python/src/typechat/_internal/ts_conversion/python_type_to_ts_nodes.py

@@ -267,7 +267,10 @@ def convert_to_type_node(py_type: object) -> TypeNode:
                         )
                         return ArrayTypeNode(AnyTypeReferenceNode)
 
-                    return ArrayTypeNode(convert_to_type_node(type_args[0]))
+                    element_type = convert_to_type_node(type_args[0])
+                    if isinstance(element_type, UnionTypeNode):
+                        return TypeReferenceNode(IdentifierNode("Array"), [element_type])
+                    return ArrayTypeNode(element_type)
 
                 return TupleTypeNode([convert_to_type_node(py_type_arg) for py_type_arg in type_args])
 

python/tests/__snapshots__/test_tuples_1/test_tuples_1.schema.d.ts

@@ -11,4 +11,8 @@ interface TupleContainer {
     arbitrary_length_4: number[];
     arbitrary_length_5: number[] | [number];
     arbitrary_length_6: number[] | [number] | [number, number];
+    union_elements: Array<number | string>;
+    nullable_elements: Array<string | null>;
+    literal_elements: Array<"yes" | "no">;
+    nested_union_elements: Array<number | string>[];
 }

python/tests/test_translator.py

@@ -47,6 +47,26 @@ def test_translator_with_immediate_pass(snapshot: Any):
 
     assert m.conversation == snapshot
 
+@dataclass
+class TupleUnionResponse:
+    values: tuple[int | str, ...]
+    optional_values: tuple[str | None, ...]
+    choices: tuple[Literal["yes", "no"], ...]
+
+
+def test_translator_with_variadic_tuple_unions():
+    model = FixedModel([
+        '{"values": [1, "two"], "optional_values": ["three", null], "choices": ["yes", "no"]}',
+    ])
+    validator = typechat.TypeChatValidator(TupleUnionResponse)
+    translator = typechat.TypeChatJsonTranslator(model, validator, TupleUnionResponse)
+
+    result = asyncio.run(translator.translate("Get tuple values."))
+
+    assert result == typechat.Success(TupleUnionResponse((1, "two"), ("three", None), ("yes", "no")))
+    assert len(model.conversation) == 2
+
+
 def test_translator_with_single_failure(snapshot: Any):
     m = FixedModel([
         '{ "a": "hello", "b": true }',

python/tests/test_tuples_1.py

@@ -1,6 +1,6 @@
 
 from dataclasses import dataclass
-from typing import Any
+from typing import Any, Literal
 
 from typechat import python_type_to_typescript_schema
 from .utilities import TypeScriptSchemaSnapshotExtension
@@ -22,6 +22,10 @@ class TupleContainer:
     arbitrary_length_4: tuple[int, ...] | tuple[float, ...]
     arbitrary_length_5: tuple[int, ...] | tuple[int]
     arbitrary_length_6: tuple[int, ...] | tuple[int] | tuple[int, int]
+    union_elements: tuple[int | str, ...]
+    nullable_elements: tuple[str | None, ...]
+    literal_elements: tuple[Literal["yes", "no"], ...]
+    nested_union_elements: tuple[tuple[int | str, ...], ...]
 
 def test_tuples_1(snapshot: Any):
-    assert python_type_to_typescript_schema(TupleContainer) == snapshot(extension_class=TypeScriptSchemaSnapshotExtension)
+    assert python_type_to_typescript_schema(TupleContainer) == snapshot(extension_class=TypeScriptSchemaSnapshotExtension)