Support inline # comments on TypedDict and dataclass fields

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

@@ -1,8 +1,10 @@
 from __future__ import annotations
 
+import ast
 from collections import OrderedDict
 import inspect
 import sys
+import textwrap
 import typing
 import typing_extensions
 from dataclasses import MISSING, Field, dataclass
@@ -152,6 +154,26 @@ class TypeScriptNodeTranslationResult:
 # }
 
 
+def _extract_inline_field_comments(py_type: type) -> dict[str, str]:
+    """Return a mapping of field name → inline # comment for a TypedDict or dataclass."""
+    try:
+        source_lines, _ = inspect.getsourcelines(py_type)
+        tree = ast.parse(textwrap.dedent("".join(source_lines)))
+        comments: dict[str, str] = {}
+        for node in ast.walk(tree):
+            if isinstance(node, ast.ClassDef):
+                for stmt in node.body:
+                    if isinstance(stmt, ast.AnnAssign) and isinstance(stmt.target, ast.Name):
+                        line = source_lines[stmt.end_lineno - 1]
+                        if "#" in line:
+                            comment = line[line.index("#") + 1:].strip().rstrip("\n")
+                            if comment:
+                                comments[stmt.target.id] = comment
+        return comments
+    except Exception:
+        return {}
+
+
 def python_type_to_typescript_nodes(root_py_type: object) -> TypeScriptNodeTranslationResult:
     # TODO: handle conflicting names
 
@@ -292,10 +314,10 @@ 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):
+    def declare_property(name: str, py_annotation: type | TypeAliasType, is_typeddict_attribute: bool, optionality_default: bool, fallback_comment: str = ""):
         """
         Declare a property for a given type.
-        If 'optionality_default' is 
+        If 'optionality_default' is
         """
         current_annotation: object = py_annotation
         origin: object
@@ -332,7 +354,7 @@ def declare_property(name: str, py_annotation: type | TypeAliasType, is_typeddic
             optional = optionality_default
 
         type_annotation = convert_to_type_node(skip_annotations(current_annotation))
-        return PropertyDeclarationNode(name, optional, comment or "", type_annotation)
+        return PropertyDeclarationNode(name, optional, comment or fallback_comment, type_annotation)
 
     def reserve_name(val: type | TypeAliasType):
         type_name = val.__name__
@@ -371,14 +393,16 @@ def declare_type(py_type: object):
                     base_attributes.setdefault(prop, set()).add(type_hint)
             bases = [convert_to_type_node(base) for base in raw_but_filtered_bases]
 
+            inline_comments = _extract_inline_field_comments(py_type)
+
             properties: list[PropertyDeclarationNode | IndexSignatureDeclarationNode] = []
             if is_typeddict(py_type):
                 for attr_name, type_hint in annotated_members.items():
                     if attribute_identical_in_all_bases(attr_name, type_hint, base_attributes):
                         continue
 
                     assume_optional = cast(TypeOfTypedDict, py_type).__total__ is False
-                    prop = declare_property(attr_name, type_hint, is_typeddict_attribute=True, optionality_default=assume_optional)
+                    prop = declare_property(attr_name, type_hint, is_typeddict_attribute=True, optionality_default=assume_optional, fallback_comment=inline_comments.get(attr_name, ""))
                     properties.append(prop)
             else:
                 # When a dataclass is created with no explicit docstring, @dataclass will
@@ -391,7 +415,7 @@ def declare_type(py_type: object):
                 for attr_name, field in cast(Dataclassish, py_type).__dataclass_fields__.items():
                     type_hint = annotated_members[attr_name]
                     optional = not(field.default is MISSING and field.default_factory is MISSING)
-                    prop = declare_property(attr_name, type_hint, is_typeddict_attribute=False, optionality_default=optional)
+                    prop = declare_property(attr_name, type_hint, is_typeddict_attribute=False, optionality_default=optional, fallback_comment=inline_comments.get(attr_name, ""))
                     properties.append(prop)
 
             reserve_name(py_type)

python/tests/__snapshots__/test_inline_comments/test_dataclass_inline_comments.schema.d.ts

@@ -0,0 +1,9 @@
+// Entry point is: 'Box'
+
+interface Box {
+    // width in pixels
+    width: number;
+    // height in pixels
+    height: number;
+    label?: string;
+}

python/tests/__snapshots__/test_inline_comments/test_doc_takes_priority_over_inline.schema.d.ts

@@ -0,0 +1,9 @@
+// Entry point is: 'Mixed'
+
+interface Mixed {
+    // from Doc
+    labeled: string;
+    // just an inline comment
+    plain: string;
+    silent: number;
+}

python/tests/__snapshots__/test_inline_comments/test_typeddict_inline_comments.schema.d.ts

@@ -0,0 +1,8 @@
+// Entry point is: 'Point'
+
+interface Point {
+    // X-coordinate
+    x: number;
+    // Y-coordinate
+    y: number;
+}

python/tests/test_inline_comments.py

@@ -0,0 +1,36 @@
+from dataclasses import dataclass
+from typing import Annotated
+from typing_extensions import Any, TypedDict, Doc
+
+from typechat import python_type_to_typescript_schema
+from .utilities import TypeScriptSchemaSnapshotExtension
+
+
+class Point(TypedDict):
+    x: int  # X-coordinate
+    y: int  # Y-coordinate
+
+
+class Mixed(TypedDict):
+    labeled: Annotated[str, Doc("from Doc")]  # inline comment (Doc takes priority)
+    plain: str  # just an inline comment
+    silent: float
+
+
+@dataclass
+class Box:
+    width: float  # width in pixels
+    height: float  # height in pixels
+    label: str = ""
+
+
+def test_typeddict_inline_comments(snapshot: Any):
+    assert python_type_to_typescript_schema(Point) == snapshot(extension_class=TypeScriptSchemaSnapshotExtension)
+
+
+def test_doc_takes_priority_over_inline(snapshot: Any):
+    assert python_type_to_typescript_schema(Mixed) == snapshot(extension_class=TypeScriptSchemaSnapshotExtension)
+
+
+def test_dataclass_inline_comments(snapshot: Any):
+    assert python_type_to_typescript_schema(Box) == snapshot(extension_class=TypeScriptSchemaSnapshotExtension)