Add support for Python tuple types (#193)

DanielRosenwasser · web-flow · commit 42d780edc89b · 2024-02-27T13:37:42.000-08:00
* Add test for tuples and update snapshots.

* Add support for tuple types.

* Update name of TypedDict to play better with upcoming Pyright versions.

* Add a test for error cases in tuples.

* Update snapshots.
diff --git a/python/src/typechat/_internal/ts_conversion/python_type_to_ts_nodes.py b/python/src/typechat/_internal/ts_conversion/python_type_to_ts_nodes.py
@@ -48,6 +48,7 @@
     StringTypeReferenceNode,
     ThisTypeReferenceNode,
     TopLevelDeclarationNode,
+    TupleTypeNode,
     TypeAliasDeclarationNode,
     TypeNode,
     TypeParameterDeclarationNode,
@@ -225,8 +226,30 @@ def convert_to_type_node(py_type: object) -> TypeNode:
                 if key_type_arg is not NumberTypeReferenceNode:
                     key_type_arg = StringTypeReferenceNode
                 return TypeReferenceNode(IdentifierNode("Record"), [key_type_arg, value_type_arg])
-
-            # TODO: tuple
+            
+            if origin is tuple:
+                # Note that when the type is `tuple[()]`,
+                # `type_args` will be an empty tuple.
+                # Which is nice, because we don't have to special-case anything!
+                type_args = get_args(py_type)
+
+                if Ellipsis in type_args:
+                    if len(type_args) != 2:
+                        errors.append(
+                            f"The tuple type '{py_type}' is ill-formed. Tuples with an ellipsis can only take the form 'tuple[SomeType, ...]'."
+                        )
+                        return ArrayTypeNode(AnyTypeReferenceNode)
+
+                    ellipsis_index = type_args.index(Ellipsis)
+                    if ellipsis_index != 1:
+                        errors.append(
+                            f"The tuple type '{py_type}' is ill-formed because the ellipsis (...) cannot be the first element."
+                        )
+                        return ArrayTypeNode(AnyTypeReferenceNode)
+                        
+                    return ArrayTypeNode(convert_to_type_node(type_args[0]))
+                    
+                return TupleTypeNode([convert_to_type_node(py_type_arg) for py_type_arg in type_args])
 
             if origin is Union or origin is UnionType:
                 type_node = [convert_to_type_node(py_type_arg) for py_type_arg in get_args(py_type)]
@@ -249,7 +272,7 @@ def convert_to_type_node(py_type: object) -> TypeNode:
 
         errors.append(f"'{py_type}' cannot be used as a type annotation.")
         return AnyTypeReferenceNode
-
+    
     def declare_property(name: str, py_annotation: type | TypeAliasType, is_typeddict_attribute: bool, optionality_default: bool):
         """
         Declare a property for a given type.
diff --git a/python/src/typechat/_internal/ts_conversion/ts_node_to_string.py b/python/src/typechat/_internal/ts_conversion/ts_node_to_string.py
@@ -10,6 +10,7 @@
     NullTypeReferenceNode,
     PropertyDeclarationNode,
     TopLevelDeclarationNode,
+    TupleTypeNode,
     TypeAliasDeclarationNode,
     TypeNode,
     TypeReferenceNode,
@@ -38,6 +39,8 @@ def ts_type_to_str(type_node: TypeNode) -> str:
             # if type(element_type) is UnionTypeNode:
             #     return f"Array<{ts_type_to_str(element_type)}>"
             return f"{ts_type_to_str(element_type)}[]"
+        case TupleTypeNode(element_types):
+            return f"[{', '.join([ts_type_to_str(element_type) for element_type in element_types])}]"
         case UnionTypeNode(types):
             # Remove duplicates, but try to preserve order of types,
             # and put null at the end if it's present.
diff --git a/python/src/typechat/_internal/ts_conversion/ts_type_nodes.py b/python/src/typechat/_internal/ts_conversion/ts_type_nodes.py
@@ -3,7 +3,7 @@
 from dataclasses import dataclass
 from typing_extensions import TypeAlias
 
-TypeNode: TypeAlias = "TypeReferenceNode | UnionTypeNode | LiteralTypeNode | ArrayTypeNode"
+TypeNode: TypeAlias = "TypeReferenceNode | UnionTypeNode | LiteralTypeNode | ArrayTypeNode | TupleTypeNode"
 
 @dataclass
 class IdentifierNode:
@@ -31,6 +31,10 @@ class LiteralTypeNode:
 class ArrayTypeNode:
     element_type: TypeNode
 
+@dataclass
+class TupleTypeNode:
+    element_types: list[TypeNode]
+
 @dataclass
 class InterfaceDeclarationNode:
     name: str
diff --git a/python/tests/__snapshots__/test_generic_alias_1.ambr b/python/tests/__snapshots__/test_generic_alias_1.ambr
diff --git a/python/tests/__snapshots__/test_tuple_errors_1.ambr b/python/tests/__snapshots__/test_tuple_errors_1.ambr
@@ -0,0 +1,4 @@
+# serializer version: 1
+# name: test_tuples_2
+  TypeScriptSchemaConversionResult(typescript_schema_str='interface TupleContainer {\n    "empty_tuples_args_1": [any, any];\n    "empty_tuples_args_2": any[];\n    "arbitrary_length_1": any[];\n    "arbitrary_length_2": any[];\n    "arbitrary_length_3": any[];\n    "arbitrary_length_4": any[];\n    "arbitrary_length_5": any[];\n    "arbitrary_length_6": any[];\n}\n', typescript_type_reference='TupleContainer', errors=["'()' cannot be used as a type annotation.", "'()' cannot be used as a type annotation.", "'()' cannot be used as a type annotation.", "The tuple type 'tuple[...]' is ill-formed. Tuples with an ellipsis can only take the form 'tuple[SomeType, ...]'.", "The tuple type 'tuple[int, int, ...]' is ill-formed. Tuples with an ellipsis can only take the form 'tuple[SomeType, ...]'.", "The tuple type 'tuple[..., int]' is ill-formed because the ellipsis (...) cannot be the first element.", "The tuple type 'tuple[..., ...]' is ill-formed because the ellipsis (...) cannot be the first element.", "The tuple type 'tuple[int, ..., int]' is ill-formed. Tuples with an ellipsis can only take the form 'tuple[SomeType, ...]'.", "The tuple type 'tuple[int, ..., int, ...]' is ill-formed. Tuples with an ellipsis can only take the form 'tuple[SomeType, ...]'."])
+# ---
diff --git a/python/tests/__snapshots__/test_tuples_1.ambr b/python/tests/__snapshots__/test_tuples_1.ambr
@@ -0,0 +1,4 @@
+# serializer version: 1
+# name: test_tuples_1
+  TypeScriptSchemaConversionResult(typescript_schema_str='interface TupleContainer {\n    "empty_tuple": [];\n    "tuple_1": [number];\n    "tuple_2": [number, string];\n    "tuple_3": [number, string];\n    "arbitrary_length_1": number[];\n    "arbitrary_length_2": number[];\n    "arbitrary_length_3": number[];\n    "arbitrary_length_4": number[];\n    "arbitrary_length_5": number[] | [number];\n    "arbitrary_length_6": number[] | [number] | [number, number];\n}\n', typescript_type_reference='TupleContainer', errors=[])
+# ---
diff --git a/python/tests/test_hello_world.py b/python/tests/test_hello_world.py
@@ -19,12 +19,12 @@ class D(C[str], total=False):
     y: Required[Annotated[bool | None, "This comes from string metadata\nwithin an Annotated hint."]]
     z: Optional[list[int]]
     other: IndirectC
-    non_class: "nonclass"
+    non_class: "NonClass"
 
     multiple_metadata: Annotated[str, None, str, "This comes from later metadata.", int]
 
 
-nonclass = TypedDict("NonClass", {"a": int, "my-dict": dict[str, int]})
+NonClass = TypedDict("NonClass", {"a": int, "my-dict": dict[str, int]})
 
 
 class E(C[str]):
diff --git a/python/tests/test_tuple_errors_1.py b/python/tests/test_tuple_errors_1.py
@@ -0,0 +1,23 @@
+
+from dataclasses import dataclass
+from typing import Any
+
+from typechat import python_type_to_typescript_schema
+
+
+@dataclass
+class TupleContainer:
+    empty_tuples_args_1: tuple[(), ()] # type: ignore
+    empty_tuples_args_2: tuple[(), ...] # type: ignore
+
+    # Arbitrary-length tuples have exactly two type arguments – the type and an ellipsis.
+    # Any other tuple form that uses an ellipsis is invalid.
+    arbitrary_length_1: tuple[...] # type: ignore
+    arbitrary_length_2: tuple[int, int, ...] # type: ignore
+    arbitrary_length_3: tuple[..., int] # type: ignore
+    arbitrary_length_4: tuple[..., ...] # type: ignore
+    arbitrary_length_5: tuple[int, ..., int] # type: ignore
+    arbitrary_length_6: tuple[int, ..., int, ...] # type: ignore
+
+def test_tuples_2(snapshot: Any):
+    assert python_type_to_typescript_schema(TupleContainer) == snapshot
diff --git a/python/tests/test_tuples_1.py b/python/tests/test_tuples_1.py
@@ -0,0 +1,27 @@
+
+from dataclasses import dataclass
+from typing import Any
+
+from typechat import python_type_to_typescript_schema
+
+
+@dataclass
+class TupleContainer:
+    # The empty tuple can be annotated as tuple[()].
+    empty_tuple: tuple[()]
+
+    tuple_1: tuple[int]
+    tuple_2: tuple[int, str]
+    tuple_3: tuple[int, str] | tuple[float, str]
+
+
+    # Arbitrary-length homogeneous tuples can be expressed using one type and an ellipsis, for example tuple[int, ...].
+    arbitrary_length_1: tuple[int, ...]
+    arbitrary_length_2: tuple[int, ...] | list[int]
+    arbitrary_length_3: tuple[int, ...] | tuple[int, ...]
+    arbitrary_length_4: tuple[int, ...] | tuple[float, ...]
+    arbitrary_length_5: tuple[int, ...] | tuple[int]
+    arbitrary_length_6: tuple[int, ...] | tuple[int] | tuple[int, int]
+
+def test_tuples_1(snapshot: Any):
+    assert python_type_to_typescript_schema(TupleContainer) == snapshot
