Implement closed support in stubgen

Author: alicederyn

mypy/stubgen.py

@@ -1094,9 +1094,13 @@ def process_typeddict(self, lvalue: NameExpr, rvalue: CallExpr) -> None:
         if not isinstance(rvalue.args[0], StrExpr):
             self.annotate_as_incomplete(lvalue)
             return
+        if len(rvalue.args) > 2 and rvalue.arg_kinds[2] != ARG_NAMED:
+            self.annotate_as_incomplete(lvalue)
+            return
 
         items: list[tuple[str, Expression]] = []
         total: Expression | None = None
+        closed: Expression | None = None
         if len(rvalue.args) > 1 and rvalue.arg_kinds[1] == ARG_POS:
             if not isinstance(rvalue.args[1], DictExpr):
                 self.annotate_as_incomplete(lvalue)
@@ -1106,11 +1110,14 @@ def process_typeddict(self, lvalue: NameExpr, rvalue: CallExpr) -> None:
                     self.annotate_as_incomplete(lvalue)
                     return
                 items.append((attr_name.value, attr_type))
-            if len(rvalue.args) > 2:
-                if rvalue.arg_kinds[2] != ARG_NAMED or rvalue.arg_names[2] != "total":
+            for arg_name, arg in zip(rvalue.arg_names[2:], rvalue.args[2:]):
+                if arg_name == "total":
+                    total = arg
+                elif arg_name == "closed":
+                    closed = arg
+                else:
                     self.annotate_as_incomplete(lvalue)
                     return
-                total = rvalue.args[2]
         else:
             for arg_name, arg in zip(rvalue.arg_names[1:], rvalue.args[1:]):
                 if not isinstance(arg_name, str):
@@ -1130,6 +1137,8 @@ def process_typeddict(self, lvalue: NameExpr, rvalue: CallExpr) -> None:
             # TODO: Add support for generic TypedDicts. Requires `Generic` as base class.
             if total is not None:
                 bases += f", total={total.accept(p)}"
+            if closed is not None:
+                bases += f", closed={closed.accept(p)}"
             class_def = f"{self._indent}class {lvalue.name}({bases}):"
             if len(items) == 0:
                 self.add(f"{class_def} ...\n")

test-data/unit/stubgen.test

@@ -3822,16 +3822,26 @@ 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)
+D1 = typing.TypedDict('D1', {'a': int, 'b': str})
+D2 = typing.TypedDict('D2', {'a': int, 'b': str}, total=False)
+D3 = typing.TypedDict('D3', {'a': int, 'b': str}, closed=True)
+D4 = typing.TypedDict('D4', {'a': int, 'b': str}, closed=True, total=False)
 [out]
 from typing_extensions import TypedDict
 
-class X(TypedDict):
+class D1(TypedDict):
     a: int
     b: str
 
-class Y(TypedDict, total=False):
+class D2(TypedDict, total=False):
+    a: int
+    b: str
+
+class D3(TypedDict, closed=True):
+    a: int
+    b: str
+
+class D4(TypedDict, total=False, closed=True):
     a: int
     b: str