PEP 655: Support Required[] inside TypedDict (#10370)

Adds support for the Required[] syntax inside TypedDicts as specified by 
PEP 655 (draft). NotRequired[] is also supported.

Co-authored-by: 97littleleaf11 <[email protected]>

Author: davidfstr

mypy/semanal.py

@@ -5127,6 +5127,7 @@ def type_analyzer(self, *,
                       allow_tuple_literal: bool = False,
                       allow_unbound_tvars: bool = False,
                       allow_placeholder: bool = False,
+                      allow_required: bool = False,
                       report_invalid_types: bool = True) -> TypeAnalyser:
         if tvar_scope is None:
             tvar_scope = self.tvar_scope
@@ -5138,8 +5139,9 @@ def type_analyzer(self, *,
                             allow_unbound_tvars=allow_unbound_tvars,
                             allow_tuple_literal=allow_tuple_literal,
                             report_invalid_types=report_invalid_types,
-                            allow_new_syntax=self.is_stub_file,
-                            allow_placeholder=allow_placeholder)
+                            allow_placeholder=allow_placeholder,
+                            allow_required=allow_required,
+                            allow_new_syntax=self.is_stub_file)
         tpan.in_dynamic_func = bool(self.function_stack and self.function_stack[-1].is_dynamic())
         tpan.global_scope = not self.type and not self.function_stack
         return tpan
@@ -5153,6 +5155,7 @@ def anal_type(self,
                   allow_tuple_literal: bool = False,
                   allow_unbound_tvars: bool = False,
                   allow_placeholder: bool = False,
+                  allow_required: bool = False,
                   report_invalid_types: bool = True,
                   third_pass: bool = False) -> Optional[Type]:
         """Semantically analyze a type.
@@ -5179,6 +5182,7 @@ def anal_type(self,
                                allow_unbound_tvars=allow_unbound_tvars,
                                allow_tuple_literal=allow_tuple_literal,
                                allow_placeholder=allow_placeholder,
+                               allow_required=allow_required,
                                report_invalid_types=report_invalid_types)
         tag = self.track_incomplete_refs()
         typ = typ.accept(a)

mypy/semanal_shared.py

@@ -119,6 +119,7 @@ def anal_type(self, t: Type, *,
                   tvar_scope: Optional[TypeVarLikeScope] = None,
                   allow_tuple_literal: bool = False,
                   allow_unbound_tvars: bool = False,
+                  allow_required: bool = False,
                   report_invalid_types: bool = True) -> Optional[Type]:
         raise NotImplementedError
 

mypy/semanal_typeddict.py

@@ -4,7 +4,9 @@
 from typing import Optional, List, Set, Tuple
 from typing_extensions import Final
 
-from mypy.types import Type, AnyType, TypeOfAny, TypedDictType, TPDICT_NAMES
+from mypy.types import (
+    Type, AnyType, TypeOfAny, TypedDictType, TPDICT_NAMES, RequiredType,
+)
 from mypy.nodes import (
     CallExpr, TypedDictExpr, Expression, NameExpr, Context, StrExpr, BytesExpr, UnicodeExpr,
     ClassDef, RefExpr, TypeInfo, AssignmentStmt, PassStmt, ExpressionStmt, EllipsisExpr, TempNode,
@@ -161,7 +163,7 @@ def analyze_typeddict_classdef_fields(
                 if stmt.type is None:
                     types.append(AnyType(TypeOfAny.unannotated))
                 else:
-                    analyzed = self.api.anal_type(stmt.type)
+                    analyzed = self.api.anal_type(stmt.type, allow_required=True)
                     if analyzed is None:
                         return None, [], set()  # Need to defer
                     types.append(analyzed)
@@ -177,7 +179,22 @@ def analyze_typeddict_classdef_fields(
             if total is None:
                 self.fail('Value of "total" must be True or False', defn)
                 total = True
-        required_keys = set(fields) if total else set()
+        required_keys = {
+            field
+            for (field, t) in zip(fields, types)
+            if (total or (
+                isinstance(t, RequiredType) and  # type: ignore[misc]
+                t.required
+            )) and not (
+                isinstance(t, RequiredType) and  # type: ignore[misc]
+                not t.required
+            )
+        }
+        types = [  # unwrap Required[T] to just T
+            t.item if isinstance(t, RequiredType) else t  # type: ignore[misc]
+            for t in types
+        ]
+
         return fields, types, required_keys
 
     def check_typeddict(self,
@@ -221,7 +238,21 @@ def check_typeddict(self,
             if name != var_name or is_func_scope:
                 # Give it a unique name derived from the line number.
                 name += '@' + str(call.line)
-            required_keys = set(items) if total else set()
+            required_keys = {
+                field
+                for (field, t) in zip(items, types)
+                if (total or (
+                    isinstance(t, RequiredType) and  # type: ignore[misc]
+                    t.required
+                )) and not (
+                    isinstance(t, RequiredType) and  # type: ignore[misc]
+                    not t.required
+                )
+            }
+            types = [  # unwrap Required[T] to just T
+                t.item if isinstance(t, RequiredType) else t  # type: ignore[misc]
+                for t in types
+            ]
             info = self.build_typeddict_typeinfo(name, items, types, required_keys, call.line)
             info.line = node.line
             # Store generated TypeInfo under both names, see semanal_namedtuple for more details.
@@ -316,7 +347,7 @@ def parse_typeddict_fields_with_types(
                 else:
                     self.fail_typeddict_arg('Invalid field type', field_type_expr)
                 return [], [], False
-            analyzed = self.api.anal_type(type)
+            analyzed = self.api.anal_type(type, allow_required=True)
             if analyzed is None:
                 return None
             types.append(analyzed)

mypy/typeanal.py

@@ -16,7 +16,7 @@
     CallableType, NoneType, ErasedType, DeletedType, TypeList, TypeVarType, SyntheticTypeVisitor,
     StarType, PartialType, EllipsisType, UninhabitedType, TypeType, CallableArgument,
     TypeQuery, union_items, TypeOfAny, LiteralType, RawExpressionType,
-    PlaceholderType, Overloaded, get_proper_type, TypeAliasType,
+    PlaceholderType, Overloaded, get_proper_type, TypeAliasType, RequiredType,
     TypeVarLikeType, ParamSpecType, ParamSpecFlavor, callable_with_ellipsis
 )
 
@@ -130,6 +130,7 @@ def __init__(self,
                  allow_new_syntax: bool = False,
                  allow_unbound_tvars: bool = False,
                  allow_placeholder: bool = False,
+                 allow_required: bool = False,
                  report_invalid_types: bool = True) -> None:
         self.api = api
         self.lookup_qualified = api.lookup_qualified
@@ -149,6 +150,8 @@ def __init__(self,
         self.allow_unbound_tvars = allow_unbound_tvars or defining_alias
         # If false, record incomplete ref if we generate PlaceholderType.
         self.allow_placeholder = allow_placeholder
+        # Are we in a context where Required[] is allowed?
+        self.allow_required = allow_required
         # Should we report an error whenever we encounter a RawExpressionType outside
         # of a Literal context: e.g. whenever we encounter an invalid type? Normally,
         # we want to report an error, but the caller may want to do more specialized
@@ -357,6 +360,22 @@ def try_analyze_special_unbound_type(self, t: UnboundType, fullname: str) -> Opt
                           " and at least one annotation", t)
                 return AnyType(TypeOfAny.from_error)
             return self.anal_type(t.args[0])
+        elif fullname in ('typing_extensions.Required', 'typing.Required'):
+            if not self.allow_required:
+                self.fail("Required[] can be only used in a TypedDict definition", t)
+                return AnyType(TypeOfAny.from_error)
+            if len(t.args) != 1:
+                self.fail("Required[] must have exactly one type argument", t)
+                return AnyType(TypeOfAny.from_error)
+            return RequiredType(self.anal_type(t.args[0]), required=True)
+        elif fullname in ('typing_extensions.NotRequired', 'typing.NotRequired'):
+            if not self.allow_required:
+                self.fail("NotRequired[] can be only used in a TypedDict definition", t)
+                return AnyType(TypeOfAny.from_error)
+            if len(t.args) != 1:
+                self.fail("NotRequired[] must have exactly one type argument", t)
+                return AnyType(TypeOfAny.from_error)
+            return RequiredType(self.anal_type(t.args[0]), required=False)
         elif self.anal_type_guard_arg(t, fullname) is not None:
             # In most contexts, TypeGuard[...] acts as an alias for bool (ignoring its args)
             return self.named_type('builtins.bool')
@@ -995,11 +1014,14 @@ def anal_array(self,
     def anal_type(self, t: Type, nested: bool = True, *, allow_param_spec: bool = False) -> Type:
         if nested:
             self.nesting_level += 1
+        old_allow_required = self.allow_required
+        self.allow_required = False
         try:
             analyzed = t.accept(self)
         finally:
             if nested:
                 self.nesting_level -= 1
+            self.allow_required = old_allow_required
         if (not allow_param_spec
                 and isinstance(analyzed, ParamSpecType)
                 and analyzed.flavor == ParamSpecFlavor.BARE):

mypy/types.py

@@ -289,7 +289,7 @@ def copy_modified(self, *,
 
 
 class TypeGuardedType(Type):
-    """Only used by find_instance_check() etc."""
+    """Only used by find_isinstance_check() etc."""
 
     __slots__ = ('type_guard',)
 
@@ -301,6 +301,21 @@ def __repr__(self) -> str:
         return "TypeGuard({})".format(self.type_guard)
 
 
+class RequiredType(Type):
+    """Required[T] or NotRequired[T]. Only usable at top-level of a TypedDict definition."""
+
+    def __init__(self, item: Type, *, required: bool) -> None:
+        super().__init__(line=item.line, column=item.column)
+        self.item = item
+        self.required = required
+
+    def __repr__(self) -> str:
+        if self.required:
+            return "Required[{}]".format(self.item)
+        else:
+            return "NotRequired[{}]".format(self.item)
+
+
 class ProperType(Type):
     """Not a type alias.
 

test-data/unit/check-typeddict.test

@@ -2146,6 +2146,132 @@ Foo = TypedDict('Foo', {'camelCaseKey': str})
 value: Foo = {}  # E: Missing key "camelCaseKey" for TypedDict "Foo"
 [builtins fixtures/dict.pyi]
 
+-- Required[]
+
+[case testDoesRecognizeRequiredInTypedDictWithClass]
+from typing import TypedDict
+from typing import Required
+class Movie(TypedDict, total=False):
+    title: Required[str]
+    year: int
+m = Movie(title='The Matrix')
+m = Movie()  # E: Missing key "title" for TypedDict "Movie"
+[typing fixtures/typing-typeddict.pyi]
+
+[case testDoesRecognizeRequiredInTypedDictWithAssignment]
+from typing import TypedDict
+from typing import Required
+Movie = TypedDict('Movie', {
+    'title': Required[str],
+    'year': int,
+}, total=False)
+m = Movie(title='The Matrix')
+m = Movie()  # E: Missing key "title" for TypedDict "Movie"
+[typing fixtures/typing-typeddict.pyi]
+
+[case testDoesDisallowRequiredOutsideOfTypedDict]
+from typing import Required
+x: Required[int] = 42  # E: Required[] can be only used in a TypedDict definition
+[typing fixtures/typing-typeddict.pyi]
+
+[case testDoesOnlyAllowRequiredInsideTypedDictAtTopLevel]
+from typing import TypedDict
+from typing import Union
+from typing import Required
+Movie = TypedDict('Movie', {
+    'title': Union[
+        Required[str],  # E: Required[] can be only used in a TypedDict definition
+        bytes
+    ],
+    'year': int,
+}, total=False)
+[typing fixtures/typing-typeddict.pyi]
+
+[case testDoesDisallowRequiredInsideRequired]
+from typing import TypedDict
+from typing import Union
+from typing import Required
+Movie = TypedDict('Movie', {
+    'title': Required[Union[
+        Required[str],  # E: Required[] can be only used in a TypedDict definition
+        bytes
+    ]],
+    'year': int,
+}, total=False)
+[typing fixtures/typing-typeddict.pyi]
+
+[case testRequiredOnlyAllowsOneItem]
+from typing import TypedDict
+from typing import Required
+class Movie(TypedDict, total=False):
+    title: Required[str, bytes]  # E: Required[] must have exactly one type argument
+    year: int
+[typing fixtures/typing-typeddict.pyi]
+
+
+-- NotRequired[]
+
+[case testDoesRecognizeNotRequiredInTypedDictWithClass]
+from typing import TypedDict
+from typing import NotRequired
+class Movie(TypedDict):
+    title: str
+    year: NotRequired[int]
+m = Movie(title='The Matrix')
+m = Movie()  # E: Missing key "title" for TypedDict "Movie"
+[typing fixtures/typing-typeddict.pyi]
+
+[case testDoesRecognizeNotRequiredInTypedDictWithAssignment]
+from typing import TypedDict
+from typing import NotRequired
+Movie = TypedDict('Movie', {
+    'title': str,
+    'year': NotRequired[int],
+})
+m = Movie(title='The Matrix')
+m = Movie()  # E: Missing key "title" for TypedDict "Movie"
+[typing fixtures/typing-typeddict.pyi]
+
+[case testDoesDisallowNotRequiredOutsideOfTypedDict]
+from typing import NotRequired
+x: NotRequired[int] = 42  # E: NotRequired[] can be only used in a TypedDict definition
+[typing fixtures/typing-typeddict.pyi]
+
+[case testDoesOnlyAllowNotRequiredInsideTypedDictAtTopLevel]
+from typing import TypedDict
+from typing import Union
+from typing import NotRequired
+Movie = TypedDict('Movie', {
+    'title': Union[
+        NotRequired[str],  # E: NotRequired[] can be only used in a TypedDict definition
+        bytes
+    ],
+    'year': int,
+})
+[typing fixtures/typing-typeddict.pyi]
+
+[case testDoesDisallowNotRequiredInsideNotRequired]
+from typing import TypedDict
+from typing import Union
+from typing import NotRequired
+Movie = TypedDict('Movie', {
+    'title': NotRequired[Union[
+        NotRequired[str],  # E: NotRequired[] can be only used in a TypedDict definition
+        bytes
+    ]],
+    'year': int,
+})
+[typing fixtures/typing-typeddict.pyi]
+
+[case testNotRequiredOnlyAllowsOneItem]
+from typing import TypedDict
+from typing import NotRequired
+class Movie(TypedDict):
+    title: NotRequired[str, bytes]  # E: NotRequired[] must have exactly one type argument
+    year: int
+[typing fixtures/typing-typeddict.pyi]
+
+-- Union dunders
 
 [case testTypedDictUnionGetItem]
 from typing import TypedDict, Union

test-data/unit/fixtures/typing-typeddict.pyi

@@ -23,6 +23,8 @@ Final = 0
 Literal = 0
 TypedDict = 0
 NoReturn = 0
+Required = 0
+NotRequired = 0
 
 T = TypeVar('T')
 T_co = TypeVar('T_co', covariant=True)