stubgen: Support TypedDict alternative syntax

mypy/stubgen.py

@@ -43,6 +43,7 @@
 
 import argparse
 import glob
+import keyword
 import os
 import os.path
 import sys
@@ -80,6 +81,7 @@
     ClassDef,
     ComparisonExpr,
     Decorator,
+    DictExpr,
     EllipsisExpr,
     Expression,
     FloatExpr,
@@ -126,6 +128,7 @@
 from mypy.traverser import all_yield_expressions, has_return_statement, has_yield_expression
 from mypy.types import (
     OVERLOAD_NAMES,
+    TPDICT_NAMES,
     AnyType,
     CallableType,
     Instance,
@@ -405,6 +408,14 @@ def visit_tuple_expr(self, node: TupleExpr) -> str:
     def visit_list_expr(self, node: ListExpr) -> str:
         return f"[{', '.join(n.accept(self) for n in node.items)}]"
 
+    def visit_dict_expr(self, o: DictExpr) -> str:
+        dict_items = []
+        for key, value in o.items:
+            # This is currently only used for TypedDict where all keys are strings.
+            assert isinstance(key, StrExpr)
+            dict_items.append(f"{key.accept(self)}: {value.accept(self)}")
+        return f"{{{', '.join(dict_items)}}}"
+
     def visit_ellipsis(self, node: EllipsisExpr) -> str:
         return "..."
 
@@ -623,7 +634,7 @@ def __init__(
         # Disable implicit exports of package-internal imports?
         self.export_less = export_less
         # Add imports that could be implicitly generated
-        self.import_tracker.add_import_from("typing", [("NamedTuple", None)])
+        self.import_tracker.add_import_from("typing", [("NamedTuple", None), ("TypedDict", None)])
         # Names in __all__ are required
         for name in _all_ or ():
             if name not in IGNORED_DUNDERS:
@@ -1003,6 +1014,10 @@ def visit_assignment_stmt(self, o: AssignmentStmt) -> None:
                 assert isinstance(o.rvalue, CallExpr)
                 self.process_namedtuple(lvalue, o.rvalue)
                 continue
+            if isinstance(lvalue, NameExpr) and self.is_typeddict(o.rvalue):
+                assert isinstance(o.rvalue, CallExpr)
+                self.process_typeddict(lvalue, o.rvalue)
+                continue
             if (
                 isinstance(lvalue, NameExpr)
                 and not self.is_private_name(lvalue.name)
@@ -1071,6 +1086,53 @@ def process_namedtuple(self, lvalue: NameExpr, rvalue: CallExpr) -> None:
                 self.add(f"{self._indent}    {item}: Incomplete\n")
         self._state = CLASS
 
+    def is_typeddict(self, expr: Expression) -> bool:
+        if not isinstance(expr, CallExpr):
+            return False
+        callee = expr.callee
+        return (
+            isinstance(callee, NameExpr) and self.refers_to_fullname(callee.name, TPDICT_NAMES)
+        ) or (
+            isinstance(callee, MemberExpr)
+            and isinstance(callee.expr, NameExpr)
+            and f"{callee.expr.name}.{callee.name}" in TPDICT_NAMES
+        )
+
+    def process_typeddict(self, lvalue: NameExpr, rvalue: CallExpr) -> None:
+        if self._state != EMPTY:
+            self.add("\n")
+        if isinstance(rvalue.args[1], DictExpr):
+            items: list[tuple[str, Expression]] = []
+            for attr_name, attr_type in rvalue.args[1].items:
+                if not isinstance(attr_name, StrExpr):
+                    self.add(f"{self._indent}{lvalue.name}: Incomplete")
+                    self.import_tracker.require_name("Incomplete")
+                    return
+                items.append((attr_name.value, attr_type))
+        else:
+            self.add(f"{self._indent}{lvalue.name}: Incomplete")
+            self.import_tracker.require_name("Incomplete")
+            return
+        self.import_tracker.require_name("TypedDict")
+        p = AliasPrinter(self)
+        if any(not key.isidentifier() or keyword.iskeyword(key) for key, _ in items):
+            # Keep the call syntax if there are non-identifier or keyword keys.
+            self.add(f"{self._indent}{lvalue.name} = {rvalue.accept(p)}\n")
+            self._state = VAR
+        else:
+            bases = "TypedDict"
+            if len(rvalue.args) > 2:
+                bases += f", total={rvalue.args[2].accept(p)}"
+            self.add(f"{self._indent}class {lvalue.name}({bases}):")
+            if len(items) == 0:
+                self.add(" ...\n")
+                self._state = EMPTY_CLASS
+            else:
+                self.add("\n")
+                for key, key_type in items:
+                    self.add(f"{self._indent}    {key}: {key_type.accept(p)}\n")
+                self._state = CLASS
+
     def is_alias_expression(self, expr: Expression, top_level: bool = True) -> bool:
         """Return True for things that look like target for an alias.
 

test-data/unit/stubgen.test

@@ -2793,3 +2793,75 @@ def f(x: str | None) -> None: ...
 a: str | int
 
 def f(x: str | None) -> None: ...
+
+[case testTypeddict]
+import typing, x
+X = typing.TypedDict('X', {'a': int, 'b': str})
+Y = typing.TypedDict('X', {'a': int, 'b': str}, total=False)
+[out]
+from typing import TypedDict
+
+class X(TypedDict):
+    a: int
+    b: str
+
+class Y(TypedDict, total=False):
+    a: int
+    b: str
+
+[case testTypeddictWithNonIdentifierOrKeywordKeys]
+from typing import TypedDict
+X = TypedDict('X', {'a-b': int, 'c': str})
+Y = TypedDict('X', {'a-b': int, 'c': str}, total=False)
+Z = TypedDict('X', {'a': int, 'in': str})
+[out]
+from typing import TypedDict
+
+X = TypedDict('X', {'a-b': int, 'c': str})
+
+Y = TypedDict('X', {'a-b': int, 'c': str}, total=False)
+
+Z = TypedDict('X', {'a': int, 'in': str})
+
+[case testEmptyTypeddict]
+import typing
+X = typing.TypedDict('X', {})
+[out]
+from typing import TypedDict
+
+class X(TypedDict): ...
+
+[case testTypeddictWithUnderscore]
+from typing import TypedDict as _TypedDict
+def f(): ...
+X = _TypedDict('X', {'a': int, 'b': str})
+def g(): ...
+[out]
+from typing import TypedDict
+
+def f() -> None: ...
+
+class X(TypedDict):
+    a: int
+    b: str
+
+def g() -> None: ...
+
+[case testNotTypeddict]
+from x import TypedDict
+import y
+X = TypedDict('X', {'a': int, 'b': str})
+Y = y.TypedDict('Y', {'a': int, 'b': str})
+[out]
+from _typeshed import Incomplete
+
+X: Incomplete
+Y: Incomplete
+
+[case testTypeddictWithWrongAttributesType]
+from typing import TypedDict
+T = TypeDict("T", {"a": int, **{"b": str, "c": bytes}})
+[out]
+from _typeshed import Incomplete
+
+T: Incomplete