From abb0bbb5c70912dd6cbb91638073cc88d4c45caf Mon Sep 17 00:00:00 2001 From: David Foster Date: Tue, 7 Dec 2021 11:16:13 -0800 Subject: [PATCH] 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 <97littleleaf11@gmail.com> --- mypy/semanal.py | 8 +- mypy/semanal_shared.py | 1 + mypy/semanal_typeddict.py | 41 +++++- mypy/typeanal.py | 24 +++- mypy/types.py | 17 ++- test-data/unit/check-typeddict.test | 126 +++++++++++++++++++ test-data/unit/fixtures/typing-typeddict.pyi | 2 + 7 files changed, 210 insertions(+), 9 deletions(-) diff --git a/mypy/semanal.py b/mypy/semanal.py index 21741f90a528..fe3151ce6cd2 100644 --- a/mypy/semanal.py +++ b/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) diff --git a/mypy/semanal_shared.py b/mypy/semanal_shared.py index 43ce1bb59214..3638d0878d8f 100644 --- a/mypy/semanal_shared.py +++ b/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 diff --git a/mypy/semanal_typeddict.py b/mypy/semanal_typeddict.py index 9611dd07e497..ffc6a7df3931 100644 --- a/mypy/semanal_typeddict.py +++ b/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) diff --git a/mypy/typeanal.py b/mypy/typeanal.py index 90f8911b1af5..f7b584eadae8 100644 --- a/mypy/typeanal.py +++ b/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): diff --git a/mypy/types.py b/mypy/types.py index d470db328374..14eefea7dd81 100644 --- a/mypy/types.py +++ b/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. diff --git a/test-data/unit/check-typeddict.test b/test-data/unit/check-typeddict.test index 776ea8fd7fc7..6e5fc88cad19 100644 --- a/test-data/unit/check-typeddict.test +++ b/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 diff --git a/test-data/unit/fixtures/typing-typeddict.pyi b/test-data/unit/fixtures/typing-typeddict.pyi index f460a7bfd167..72f500707094 100644 --- a/test-data/unit/fixtures/typing-typeddict.pyi +++ b/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)