Python: support typing.NewType in schema conversion

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

@@ -226,6 +226,11 @@ def convert_to_type_node(py_type: object) -> TypeNode:
         if py_type is Self:
             return ThisTypeReferenceNode
 
+        # A NewType has the same runtime representation as its underlying type,
+        # so we emit the underlying type (which may itself be a NewType).
+        if isinstance(py_type, (typing.NewType, typing_extensions.NewType)):
+            return convert_to_type_node(py_type.__supertype__)
+
         # TODO: consider handling bare 'tuple' (and list, etc.)
         # https://docs.python.org/3/library/typing.html#annotating-tuples
         # Using plain tuple as an annotation is equivalent to using tuple[Any, ...]:

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

@@ -0,0 +1,11 @@
+// Entry point is: 'Profile'
+
+interface Profile {
+    user_id: string;
+    count: number;
+    admin_id: string;
+    tags: string[];
+    scores: Record<string, number>;
+    optional_id: string | null;
+    maybe_count?: number;
+}

python/tests/__snapshots__/test_translator.ambr

@@ -117,6 +117,39 @@
     }),
   ])
 # ---
+# name: test_translator_with_newtype
+  list([
+    dict({
+      'kind': 'CLIENT REQUEST',
+      'payload': list([
+        dict({
+          'content': '''
+
+            You are a service that translates user requests into JSON objects of type "ExampleWithNewType" according to the following TypeScript definitions:
+            ```
+            interface ExampleWithNewType {
+                user_id: string;
+                count: number;
+            }
+
+            ```
+            The following is a user request:
+            '''
+            Get me a user.
+            '''
+            The following is the user request translated into a JSON object with 2 spaces of indentation and no properties with the value undefined:
+
+          ''',
+          'role': 'user',
+        }),
+      ]),
+    }),
+    dict({
+      'kind': 'MODEL RESPONSE',
+      'payload': '{ "user_id": "abc", "count": 1 }',
+    }),
+  ])
+# ---
 # name: test_translator_with_single_failure
   list([
     dict({

python/tests/test_newtype_1.py

@@ -0,0 +1,37 @@
+from typing import Any, NewType
+from typing_extensions import NotRequired, TypedDict
+
+import typechat
+from typechat import python_type_to_typescript_schema
+from .utilities import TypeScriptSchemaSnapshotExtension
+
+UserId = NewType("UserId", str)
+Count = NewType("Count", int)
+AdminId = NewType("AdminId", UserId)
+
+class Profile(TypedDict):
+    user_id: UserId
+    count: Count
+    admin_id: AdminId
+    tags: list[UserId]
+    scores: dict[str, Count]
+    optional_id: UserId | None
+    maybe_count: NotRequired[Count]
+
+def test_newtype_1(snapshot: Any):
+    result = python_type_to_typescript_schema(Profile)
+    assert result.errors == []
+    assert result == snapshot(extension_class=TypeScriptSchemaSnapshotExtension)
+
+def test_newtype_validation():
+    v = typechat.TypeChatValidator(Profile)
+    valid = {"user_id": "a", "count": 1, "admin_id": "b", "tags": ["c"], "scores": {"x": 2}, "optional_id": None}
+    assert isinstance(v.validate_object(valid), typechat.Success)
+    assert isinstance(v.validate_object({**valid, "user_id": 42}), typechat.Failure)
+    assert isinstance(v.validate_object({**valid, "count": "1"}), typechat.Failure)
+
+class Unsupported(TypedDict):
+    value: NewType("Weird", 1.5)  # type: ignore
+
+def test_newtype_unsupported_underlying_type_reports_error():
+    assert python_type_to_typescript_schema(Unsupported).errors != []

python/tests/test_translator.py

@@ -1,6 +1,7 @@
 
 import asyncio
 from dataclasses import dataclass
+from typing import NewType
 from typing_extensions import Any, Iterator, Literal, TypedDict, override
 import typechat
 
@@ -93,3 +94,21 @@ def test_translator_with_single_failure_and_list_preamble_1(snapshot: Any):
 
     assert m.conversation == snapshot
 
+UserId = NewType("UserId", str)
+
+@dataclass
+class ExampleWithNewType:
+    user_id: UserId
+    count: int
+
+v_newtype = typechat.TypeChatValidator(ExampleWithNewType)
+
+def test_translator_with_newtype(snapshot: Any):
+    m = FixedModel([
+        '{ "user_id": "abc", "count": 1 }',
+    ])
+    t = typechat.TypeChatJsonTranslator(m, v_newtype, ExampleWithNewType)
+    asyncio.run(t.translate("Get me a user."))
+
+    assert m.conversation == snapshot
+

site/src/docs/python/basic-usage.md

@@ -146,6 +146,11 @@ That would mean that to access the value of `sentiment`, we would have to write
 Note that while we used [the built-in `typing` module](https://docs.python.org/3/library/typing.html), [`typing_extensions`](https://pypi.org/project/typing-extensions/) is supported as well.
 TypeChat also understands constructs like `Annotated` and `Doc` to add comments to individual attributes.
 
+[`typing.NewType`](https://docs.python.org/3/library/typing.html#typing.NewType) (and `typing_extensions.NewType`) is also supported.
+TypeChat converts a NewType through its `__supertype__`, so the TypeScript schema uses the underlying JSON type: `NewType("UserId", str)` becomes `string`, and a NewType built on another NewType follows the chain.
+Runtime validation still uses the declared Python type; NewType itself does not add checks beyond that underlying type.
+
+
 ## Creating a Validator
 
 A validator really has two jobs generating a textual schema for language models, and making sure any data fits a given shape.