Implement support for closed TypedDicts (PEP 728) (#21382)

Implement support for the closed keyword on TypedDicts (part of [PEP
728][]).

Additionally, fix some preexisting issues that I came across while
updating the logic.

* Fixes #7435
* Fixes #7981
* Fixes #8714
* Fixes #12143
* Fixes #20401
* Partially addresses #18176

[PEP 728]: https://peps.python.org/pep-0728/

Author: alicederyn

docs/source/typed_dict.rst

@@ -81,9 +81,9 @@ arbitrarily complex types. For example, you can define nested
 ``TypedDict``\s and containers with ``TypedDict`` items.
 Unlike most other types, mypy uses structural compatibility checking
 (or structural subtyping) with ``TypedDict``\s. A ``TypedDict`` object with
-extra items is compatible with (a subtype of) a narrower
+extra items can be compatible with (a subtype of) a narrower
 ``TypedDict``, assuming item types are compatible (*totality* also affects
-subtyping, as discussed below).
+subtyping, as does *closing*, as discussed below).
 
 A ``TypedDict`` object is not a subtype of the regular ``dict[...]``
 type (and vice versa), since :py:class:`dict` allows arbitrary keys to be
@@ -276,6 +276,84 @@ vary :ref:`covariantly <variance-of-generics>`:
     m: Movie = {"name": "Jaws", "year": 1975}
     process_entry(m)  # OK
 
+You can override a read-only item with a compatible subtype, make a
+read-only item mutable, and inherit from multiple parents with compatible
+definitions:
+
+.. code-block:: python
+
+    from collections.abc import Collection, Sequence
+
+    class Competition(TypedDict):
+        hosts: ReadOnly[Collection[str]]
+        entries: ReadOnly[Sequence[Entry]]
+
+    class MovieShow(TypedDict):
+        entries: list[Movie]
+
+    class Oscars(Competition, MovieShow):
+        hosts: set[str]
+
+Defining ``hosts`` as a mutable ``set[str]`` item works as this is compatible
+with the read-only ``Collection[str]`` definition in ``Competition``.
+``entries`` will be of type ``list[Movie]``, taken from the ``MovieShow`` type,
+as it is the only non-readonly definition, and is compatible with the definition
+in ``Competition``.
+
+If an item is only defined in supertypes, and is always read-only, mypy takes
+the definition from the first parent in the inheritance order, and raises an
+error if any other parent definition is incompatible:
+
+.. code-block:: python
+
+    class NameIds(TypedDict):
+        ids: ReadOnly[Collection[str]]
+
+    class OrderedIds(TypedDict):
+        ids: ReadOnly[Sequence[int | str]]
+
+    class OrderedNameIds(NameIds, OrderedIds):
+        pass  # Error! Parent definitions incompatible
+
+In this example, the definition of ``ids`` will be taken from ``NameIds``,
+which would not be compatible with the definition in ``OrderedIds``; reordering
+the parents would not solve the problem. Instead, you will need to make a
+compatible definition explicitly:
+
+.. code-block:: python
+
+    class OrderedNameIds(NameIds, OrderedIds):
+        ids: ReadOnly[Sequence[str]]
+
+Closing
+-------
+
+You can use the ``closed`` keyword, introduced to ``TypedDict`` in Python
+3.15 (and available via ``typing_extensions.TypedDict`` in older
+versions) to prevent structural subtypes from adding extra keys to a
+type (:pep:`728`):
+
+.. code-block:: python
+
+   HasName = TypedDict("HasName", {"name": str})
+   HasOnlyName = TypedDict("HasOnlyName", {"name": str}, closed=True)
+   Movie = TypedDict("Movie", {"name": str, "year": int})
+
+   movie: Movie = {"name": "Nimona", "year": 2023}
+   has_name: HasName = movie  # OK: type is open
+   has_only_name: HasOnlyName = movie  # Error: type is closed
+
+This allows the typechecker to determine that certain operations are safe,
+when they otherwise wouldn't be due to the potential presence of unknown
+keys.
+
+The ``closed`` keyword can also be used in class-based syntax:
+
+.. code-block:: python
+
+   class HasOnlyName(TypedDict, closed=True):
+       name: str
+
 Unions of TypedDicts
 --------------------
 
@@ -289,6 +367,40 @@ need to give each TypedDict the same key where each value has a unique
 :ref:`Literal type <literal_types>`. Then, check that key to distinguish
 between your TypedDicts.
 
+Alternatively, you can implement tagged unions with single-key wrapper dictionaries:
+
+.. code-block:: python
+
+   class Book(TypedDict):
+       name: str
+       length: int
+       ...
+
+   class DVD(TypedDict):
+       name: str
+       length: int
+       ...
+
+   TaggedBook = TypedDict('TaggedBook', {'book': Book}, closed=True)
+   TaggedDVD = TypedDict('TaggedDVD', {'dvd': DVD}, closed=True)
+   type Inventory = TaggedBook | TaggedDVD
+
+   def print_length(inventory: Inventory) -> None:
+       if "book" in inventory:
+           print(inventory["book"]["length"], 'pages')
+       else:
+           print(inventory["dvd"]["length"], 'minutes')
+
+Here, the ``closed`` keyword is necessary to allow the ``if`` guard to safely
+narrow the types; without it, there could be a structural subtype of ``TaggedDVD``
+that contains a ``book`` field of arbitrary type.
+
+.. note::
+
+   Applying ``@final`` to a TypedDict is a legacy way of marking it as closed
+   for the purposes of type narrowing. It was never fully implemented and is
+   now superseded; it may be removed in future.
+
 Inline TypedDict types
 ----------------------
 

mypy/checker.py

@@ -2813,6 +2813,8 @@ def visit_class_def(self, defn: ClassDef) -> None:
                     self.check_multiple_inheritance(typ)
                 self.check_metaclass_compatibility(typ)
                 self.check_final_deletable(typ)
+                if typ.typeddict_type:
+                    self.check_typeddict_inheritance(defn)
 
             if defn.decorators:
                 sig: Type = type_object_type(defn.info)
@@ -3219,6 +3221,44 @@ def check_metaclass_compatibility(self, typ: TypeInfo) -> None:
             if explanation:
                 self.note(explanation, typ, code=codes.METACLASS)
 
+    def check_typeddict_inheritance(self, defn: ClassDef) -> None:
+        """Ensure that the final definition of a TypedDict is compatible with its base classes."""
+        assert defn.info.typeddict_type
+        td = defn.info.typeddict_type
+        data = defn.info.typeddict_data
+        if data is None or not data.ready:
+            return
+        for base, base_items in data.bases:
+            assert base.typeddict_type
+            for field_name, base_type in base_items.items():
+                field_type = td.items[field_name]
+                assert field_type
+                is_readonly = field_name in base.typeddict_type.readonly_keys
+                if is_readonly:
+                    is_compatible = is_subtype(field_type, base_type)
+                else:
+                    is_compatible = is_equivalent(field_type, base_type)
+                if not is_compatible:
+                    source = data.field_sources[field_name]
+                    if source.base is None:
+                        self.fail(
+                            f'Definition of field "{field_name}" incompatible with base class '
+                            f'"{base.name}"',
+                            source.ctx,
+                        )
+                    else:
+                        self.fail(
+                            f'Incompatible definitions of field "{field_name}" in base classes '
+                            f'"{base.name}" and "{source.base.name}"',
+                            source.ctx,
+                        )
+                        if field_name in td.readonly_keys:
+                            self.note(
+                                f'This can be resolved by redeclaring the field "{field_name}" '
+                                f"with a mutually compatible type",
+                                source.ctx,
+                            )
+
     def visit_import_from(self, node: ImportFrom) -> None:
         for name, _ in node.names:
             if (sym := self.globals.get(name)) is not None:
@@ -6393,7 +6433,8 @@ def conditional_types_for_iterable(
     ) -> tuple[Type, Type]:
         """
         Narrows the type of `iterable_type` based on the type of `item_type`.
-        For now, we only support narrowing unions of TypedDicts based on left operand being literal string(s).
+        For now, we only support narrowing unions of TypedDicts, and TypeVars with TypedDict
+        bounds, based on left operand being literal string(s).
         """
         if_types: list[Type] = []
         else_types: list[Type] = []
@@ -6407,16 +6448,30 @@ def conditional_types_for_iterable(
         item_str_literals = try_getting_str_literals_from_type(item_type)
 
         for possible_iterable_type in possible_iterable_types:
-            if item_str_literals and isinstance(possible_iterable_type, TypedDictType):
+            bound = (
+                get_proper_type(possible_iterable_type.upper_bound)
+                if isinstance(possible_iterable_type, TypeVarType)
+                else possible_iterable_type
+            )
+
+            if item_str_literals and isinstance(bound, TypedDictType):
                 for key in item_str_literals:
-                    if key in possible_iterable_type.required_keys:
+                    if key in bound.required_keys:
                         if_types.append(possible_iterable_type)
-                    elif (
-                        key in possible_iterable_type.items or not possible_iterable_type.is_final
+                    elif key in bound.items and isinstance(
+                        get_proper_type(bound.items[key]), UninhabitedType
                     ):
-                        if_types.append(possible_iterable_type)
+                        # If an item is explicitly declared uninhabited, we can exclude it from
+                        # if_types; see testOperatorContainsNarrowsTypedDicts_closed
+                        else_types.append(possible_iterable_type)
+                    elif key not in bound.items and (bound.is_closed or bound.is_final):
+                        # If an item is missing and the type is closed, we can exclude it from
+                        # if_types; see testOperatorContainsNarrowsTypedDicts_closed
+                        # We also support "final" as a legacy way of expressing "closed" in this
+                        # specific case; see testOperatorContainsNarrowsTypedDicts_final
                         else_types.append(possible_iterable_type)
                     else:
+                        if_types.append(possible_iterable_type)
                         else_types.append(possible_iterable_type)
             else:
                 if_types.append(possible_iterable_type)

mypy/checkexpr.py

@@ -818,8 +818,8 @@ def validate_typeddict_kwargs(
         result = defaultdict(list)
         # Keys that are guaranteed to be present no matter what (e.g. for all items of a union)
         always_present_keys = set()
-        # Indicates latest encountered ** unpack among items.
-        last_star_found = None
+        # Indicates latest encountered ** unpack of a non-closed type among items.
+        last_open_star_found = None
 
         for item_name_expr, item_arg in kwargs:
             if item_name_expr:
@@ -843,22 +843,30 @@ def validate_typeddict_kwargs(
                     result[literal_value] = [item_arg]
                     always_present_keys.add(literal_value)
             else:
-                last_star_found = item_arg
-                if not self.validate_star_typeddict_item(
+                is_valid, is_open = self.validate_star_typeddict_item(
                     item_arg, callee, result, always_present_keys
-                ):
+                )
+                if not is_valid:
                     return None
-        if self.chk.options.extra_checks and last_star_found is not None:
+                if is_open:
+                    last_open_star_found = item_arg
+        if self.chk.options.extra_checks and last_open_star_found is not None:
+            if callee.is_closed:
+                self.chk.fail(
+                    "Cannot unpack item that may contain extra keys into a closed TypedDict",
+                    last_open_star_found,
+                    code=codes.TYPEDDICT_ITEM,
+                )
             absent_keys = []
             for key in callee.items:
                 if key not in callee.required_keys and key not in result:
                     absent_keys.append(key)
             if absent_keys:
-                # Having an optional key not explicitly declared by a ** unpacked
+                # Having an optional key not explicitly declared by a ** unpacked open
                 # TypedDict is unsafe, it may be an (incompatible) subtype at runtime.
                 # TODO: catch the cases where a declared key is overridden by a subsequent
                 # ** item without it (and not again overridden with complete ** item).
-                self.msg.non_required_keys_absent_with_star(absent_keys, last_star_found)
+                self.msg.non_required_keys_absent_with_star(absent_keys, last_open_star_found)
         return result, always_present_keys
 
     def validate_star_typeddict_item(
@@ -867,14 +875,18 @@ def validate_star_typeddict_item(
         callee: TypedDictType,
         result: dict[str, list[Expression]],
         always_present_keys: set[str],
-    ) -> bool:
+    ) -> tuple[bool, bool]:
         """Update keys/expressions from a ** expression in TypedDict constructor.
 
-        Note `result` and `always_present_keys` are updated in place. Return true if the
-        expression `item_arg` may valid in `callee` TypedDict context.
+        Note `result` and `always_present_keys` are updated in place.
+
+        First tuple item returned is true if the expression `item_arg` may valid
+        in `callee` TypedDict context. Second tuple item returned is true if the
+        expression may contain other keys not explicitly declared.
         """
         inferred = get_proper_type(self.accept(item_arg, type_context=callee))
-        possible_tds = []
+        any_fallback = False
+        possible_tds: list[TypedDictType] = []
         if isinstance(inferred, TypedDictType):
             possible_tds = [inferred]
         elif isinstance(inferred, UnionType):
@@ -883,10 +895,14 @@ def validate_star_typeddict_item(
                     possible_tds.append(item)
                 elif not self.valid_unpack_fallback_item(item):
                     self.msg.unsupported_target_for_star_typeddict(item, item_arg)
-                    return False
+                    return False, True
+                else:
+                    any_fallback = True
         elif not self.valid_unpack_fallback_item(inferred):
             self.msg.unsupported_target_for_star_typeddict(inferred, item_arg)
-            return False
+            return False, True
+        else:
+            any_fallback = True
         all_keys: set[str] = set()
         for td in possible_tds:
             all_keys |= td.items.keys()
@@ -915,7 +931,8 @@ def validate_star_typeddict_item(
                 # If this key is not required at least in some item of a union
                 # it may not shadow previous item, so we need to type check both.
                 result[key].append(arg)
-        return True
+        all_closed = all(t.is_closed for t in possible_tds)
+        return True, any_fallback or not all_closed
 
     def valid_unpack_fallback_item(self, typ: ProperType) -> bool:
         if isinstance(typ, AnyType):

mypy/copytype.py

@@ -107,7 +107,7 @@ def visit_tuple_type(self, t: TupleType) -> ProperType:
 
     def visit_typeddict_type(self, t: TypedDictType) -> ProperType:
         return self.copy_common(
-            t, TypedDictType(t.items, t.required_keys, t.readonly_keys, t.fallback)
+            t, TypedDictType(t.items, t.required_keys, t.readonly_keys, t.is_closed, t.fallback)
         )
 
     def visit_literal_type(self, t: LiteralType) -> ProperType:

mypy/expandtype.py

@@ -336,7 +336,7 @@ def _possible_callable_kwargs(cls, repl: Parameters, dict_type: Instance) -> Pro
             return dict_type
         kwargs = {}
         required_names = set()
-        extra_items: Type = UninhabitedType()
+        extra_items: Type | None = None
         for kind, name, type in zip(repl.arg_kinds, repl.arg_names, repl.arg_types):
             if kind == ArgKind.ARG_NAMED and name is not None:
                 kwargs[name] = type
@@ -346,10 +346,11 @@ def _possible_callable_kwargs(cls, repl: Parameters, dict_type: Instance) -> Pro
                 extra_items = type
             elif not kind.is_star() and name is not None:
                 kwargs[name] = type
-        if not kwargs:
+        if not kwargs and extra_items is not None:
             return Instance(dict_type.type, [dict_type.args[0], extra_items])
-        # TODO: when PEP 728 is implemented, pass extra_items below.
-        return TypedDictType(kwargs, required_names, set(), fallback=dict_type)
+        # TODO: when PEP 728 `extra_items` is implemented, pass extra_items below.
+        is_closed = extra_items is None
+        return TypedDictType(kwargs, required_names, set(), is_closed, fallback=dict_type)
 
     def visit_type_var_tuple(self, t: TypeVarTupleType) -> Type:
         # Sometimes solver may need to expand a type variable with (a copy of) itself

mypy/exprtotype.py

@@ -277,7 +277,7 @@ def expr_to_unanalyzed_type(
                 value, options, allow_new_syntax, expr, lookup_qualified=lookup_qualified
             )
         result = TypedDictType(
-            items, set(), set(), Instance(MISSING_FALLBACK, ()), expr.line, expr.column
+            items, set(), set(), False, Instance(MISSING_FALLBACK, ()), expr.line, expr.column
         )
         result.extra_items_from = extra_items_from
         return result

mypy/fastparse.py

@@ -2143,7 +2143,7 @@ def visit_Dict(self, n: ast3.Dict) -> Type:
                     continue
                 return self.invalid_type(n)
             items[item_name.value] = self.visit(value)
-        result = TypedDictType(items, set(), set(), _dummy_fallback, n.lineno, n.col_offset)
+        result = TypedDictType(items, set(), set(), False, _dummy_fallback, n.lineno, n.col_offset)
         result.extra_items_from = extra_items_from
         return result