Support __closed__ and __extra_items__ for PEP 728.

Signed-off-by: Zixuan James Li <[email protected]>

Author: PIG208

CHANGELOG.md

@@ -1,5 +1,7 @@
 # Unreleased
 
+- Add support for PEP 728, supporting the `closed` keyword argument and the
+  special `__extra_items__` key for TypedDict. Patch by Zixuan James Li.
 - Add support for PEP 742, adding `typing_extensions.TypeIs`. Patch
   by Jelle Zijlstra.
 - Drop runtime error when a read-only `TypedDict` item overrides a mutable

doc/index.rst

@@ -394,6 +394,32 @@ Special typing primitives
       are mutable if they do not carry the :data:`ReadOnly` qualifier.
 
       .. versionadded:: 4.9.0
+    
+    The experimental ``closed`` keyword argument and the special key
+    ``"__extra_items__"`` proposed in :pep:`728` are supported.
+
+    When ``closed`` is unspecified or ``closed=False`` is given,
+    ``"__extra_items__"`` behave like a regular key. Otherwise, this becomes a
+    special key that does not show up in ``__readonly_keys__``,
+    ``__mutable_keys__``, ``__required_keys__``, ``__optional_keys``, or
+    ``__annotations__``.
+
+    For runtime introspection, two attributes can be looked at:
+
+    .. attribute:: __closed__
+
+        A boolean flag indicating the value of the keyword argument ``closed``
+        on the current ``TypedDict``.
+
+        .. versionadded:: 4.10.0
+
+    .. attribute:: __extra_items__
+
+        The type annotation of the extra items allowed on the ``TypedDict``.
+        This attribute does not appear on a TypedDict that has itself and all
+        its bases non-closed.
+
+        .. versionadded:: 4.10.0
 
    .. versionchanged:: 4.3.0
 
@@ -427,6 +453,11 @@ Special typing primitives
 
       Support for the :data:`ReadOnly` qualifier was added.
 
+   .. versionchanged:: 4.10.0
+
+      The keyword argument ``closed`` and the special key ``"__extra_items__"``
+      when ``closed=True`` is given were supported.
+
 .. class:: TypeVar(name, *constraints, bound=None, covariant=False,
                    contravariant=False, infer_variance=False, default=...)
 

src/test_typing_extensions.py

@@ -4177,6 +4177,121 @@ class AllTheThings(TypedDict):
         self.assertEqual(AllTheThings.__optional_keys__, frozenset({'c', 'd'}))
         self.assertEqual(AllTheThings.__readonly_keys__, frozenset({'a', 'b', 'c'}))
         self.assertEqual(AllTheThings.__mutable_keys__, frozenset({'d'}))
+    
+    def test_extra_keys_non_readonly(self):
+        class Base(TypedDict, closed=True):
+            __extra_items__: str
+        
+        class Child(Base):
+            a: NotRequired[int]
+        
+        self.assertEqual(Child.__required_keys__, frozenset({}))
+        self.assertEqual(Child.__optional_keys__, frozenset({'a'}))
+        self.assertEqual(Child.__readonly_keys__, frozenset({}))
+        self.assertEqual(Child.__mutable_keys__, frozenset({'a'}))
+
+    def test_extra_keys_readonly(self):
+        class Base(TypedDict, closed=True):
+            __extra_items__: ReadOnly[str]
+        
+        class Child(Base):
+            a: NotRequired[str]
+        
+        self.assertEqual(Child.__required_keys__, frozenset({}))
+        self.assertEqual(Child.__optional_keys__, frozenset({'a'}))
+        self.assertEqual(Child.__readonly_keys__, frozenset({}))
+        self.assertEqual(Child.__mutable_keys__, frozenset({'a'}))
+    
+    def test_extra_key_required(self):
+        with self.assertRaisesRegex(
+            TypeError,
+            "Special key __extra_items__ does not support Required and NotRequired"
+        ):
+            TypedDict("A", {"__extra_items__": Required[int]}, closed=True)
+
+        with self.assertRaisesRegex(
+            TypeError,
+            "Special key __extra_items__ does not support Required and NotRequired"
+        ):
+            TypedDict("A", {"__extra_items__": NotRequired[int]}, closed=True)
+    
+    def test_regular_extra_items(self):
+        class ExtraReadOnly(TypedDict):
+            __extra_items__: ReadOnly[str]
+
+        class ExtraRequired(TypedDict):
+            __extra_items__: Required[str]
+
+        class ExtraNotRequired(TypedDict):
+            __extra_items__: NotRequired[str]
+        
+        self.assertEqual(ExtraReadOnly.__required_keys__, frozenset({'__extra_items__'}))
+        self.assertEqual(ExtraReadOnly.__optional_keys__, frozenset({}))
+        self.assertEqual(ExtraReadOnly.__readonly_keys__, frozenset({'__extra_items__'}))
+        self.assertEqual(ExtraReadOnly.__mutable_keys__, frozenset({}))
+
+        self.assertEqual(ExtraRequired.__required_keys__, frozenset({'__extra_items__'}))
+        self.assertEqual(ExtraRequired.__optional_keys__, frozenset({}))
+        self.assertEqual(ExtraRequired.__readonly_keys__, frozenset({}))
+        self.assertEqual(ExtraRequired.__mutable_keys__, frozenset({'__extra_items__'}))
+
+        self.assertEqual(ExtraNotRequired.__required_keys__, frozenset({}))
+        self.assertEqual(ExtraNotRequired.__optional_keys__, frozenset({'__extra_items__'}))
+        self.assertEqual(ExtraNotRequired.__readonly_keys__, frozenset({}))
+        self.assertEqual(ExtraNotRequired.__mutable_keys__, frozenset({'__extra_items__'}))
+    
+    def test_closed_inheritance(self):
+        class Base(TypedDict, closed=True):
+            __extra_items__: ReadOnly[Union[str, None]]
+        
+        class Child(Base):
+            a: int
+            __extra_items__: int
+
+        class GrandChild(Child, closed=True):
+            __extra_items__: str
+
+        self.assertEqual(Base.__required_keys__, frozenset({}))
+        self.assertEqual(Base.__optional_keys__, frozenset({}))
+        self.assertEqual(Base.__readonly_keys__, frozenset({}))
+        self.assertEqual(Base.__mutable_keys__, frozenset({}))
+        self.assertEqual(Base.__extra_items__, ReadOnly[Union[str, None]])
+
+        self.assertEqual(Child.__required_keys__, frozenset({'a', "__extra_items__"}))
+        self.assertEqual(Child.__optional_keys__, frozenset({}))
+        self.assertEqual(Child.__readonly_keys__, frozenset({}))
+        self.assertEqual(Child.__mutable_keys__, frozenset({'a', "__extra_items__"}))
+        self.assertEqual(Child.__extra_items__, ReadOnly[Union[str, None]])
+
+        self.assertEqual(GrandChild.__required_keys__, frozenset({'a', "__extra_items__"}))
+        self.assertEqual(GrandChild.__optional_keys__, frozenset({}))
+        self.assertEqual(GrandChild.__readonly_keys__, frozenset({}))
+        self.assertEqual(GrandChild.__mutable_keys__, frozenset({'a', "__extra_items__"}))
+        self.assertEqual(GrandChild.__extra_items__, str)
+
+        self.assertEqual(Base.__annotations__, {})
+        self.assertEqual(Child.__annotations__, {"__extra_items__": int, "a": int})
+        self.assertEqual(GrandChild.__annotations__, {"__extra_items__": int, "a": int})
+
+        self.assertTrue(Base.__closed__)
+        self.assertFalse(Child.__closed__)
+        self.assertTrue(GrandChild.__closed__)
+    
+    def test_absent_extra_items(self):
+        class Base(TypedDict):
+            a: int
+        
+        class ChildA(Base, closed=True):
+            ...
+
+        class ChildB(Base, closed=True):
+            __extra_items__: None
+        
+        self.assertNotIn("__extra_items__", Base.__dict__)
+        self.assertIn("__extra_items__", ChildA.__dict__)
+        self.assertIn("__extra_items__", ChildB.__dict__)
+        self.assertEqual(ChildA.__extra_items__, Never)
+        self.assertEqual(ChildB.__extra_items__, type(None))
 
 
 class AnnotatedTests(BaseTestCase):

src/typing_extensions.py

@@ -875,7 +875,7 @@ def _get_typeddict_qualifiers(annotation_type):
                 break
 
     class _TypedDictMeta(type):
-        def __new__(cls, name, bases, ns, *, total=True):
+        def __new__(cls, name, bases, ns, *, total=True, closed=False):
             """Create new typed dict class object.
 
             This method is called when TypedDict is subclassed,
@@ -920,6 +920,7 @@ def __new__(cls, name, bases, ns, *, total=True):
             optional_keys = set()
             readonly_keys = set()
             mutable_keys = set()
+            extra_items_type = _marker
 
             for base in bases:
                 base_dict = base.__dict__
@@ -929,6 +930,20 @@ def __new__(cls, name, bases, ns, *, total=True):
                 optional_keys.update(base_dict.get('__optional_keys__', ()))
                 readonly_keys.update(base_dict.get('__readonly_keys__', ()))
                 mutable_keys.update(base_dict.get('__mutable_keys__', ()))
+                if '__extra_items__' in base_dict:
+                    extra_items_type = base_dict['__extra_items__']
+
+            if closed and extra_items_type is _marker:
+                extra_items_type = Never
+            if closed and "__extra_items__" in own_annotations:
+                annotation_type = own_annotations.pop("__extra_items__")
+                qualifiers = set(_get_typeddict_qualifiers(annotation_type))
+                if Required in qualifiers or NotRequired in qualifiers:
+                    raise TypeError(
+                        f"Special key __extra_items__ does not support"
+                        " Required and NotRequired"
+                    )
+                extra_items_type = annotation_type
 
             annotations.update(own_annotations)
             for annotation_key, annotation_type in own_annotations.items():
@@ -956,6 +971,9 @@ def __new__(cls, name, bases, ns, *, total=True):
             tp_dict.__mutable_keys__ = frozenset(mutable_keys)
             if not hasattr(tp_dict, '__total__'):
                 tp_dict.__total__ = total
+            tp_dict.__closed__ = closed
+            if extra_items_type is not _marker:
+                tp_dict.__extra_items__ = extra_items_type
             return tp_dict
 
         __call__ = dict  # static method
@@ -969,7 +987,7 @@ def __subclasscheck__(cls, other):
     _TypedDict = type.__new__(_TypedDictMeta, 'TypedDict', (), {})
 
     @_ensure_subclassable(lambda bases: (_TypedDict,))
-    def TypedDict(typename, fields=_marker, /, *, total=True, **kwargs):
+    def TypedDict(typename, fields=_marker, /, *, total=True, closed=False, **kwargs):
         """A simple typed namespace. At runtime it is equivalent to a plain dict.
 
         TypedDict creates a dictionary type such that a type checker will expect all
@@ -1050,7 +1068,7 @@ class Point2D(TypedDict):
             # Setting correct module is necessary to make typed dict classes pickleable.
             ns['__module__'] = module
 
-        td = _TypedDictMeta(typename, (), ns, total=total)
+        td = _TypedDictMeta(typename, (), ns, total=total, closed=closed)
         td.__orig_bases__ = (TypedDict,)
         return td