Allow using typeddict for more precise typing of kwargs

mypy-requirements.txt

@@ -1,4 +1,4 @@
 typing_extensions>=3.7.4
-mypy_extensions>=0.4.3,<0.5.0
+mypy_extensions>=0.4.3,<0.6.0
 typed_ast>=1.4.0,<1.5.0
 toml

mypy/semanal.py

@@ -76,7 +76,7 @@
     get_nongen_builtins, get_member_expr_fullname, REVEAL_TYPE,
     REVEAL_LOCALS, is_final_node, TypedDictExpr, type_aliases_source_versions,
     EnumCallExpr, RUNTIME_PROTOCOL_DECOS, FakeExpression, Statement, AssignmentExpr,
-    ParamSpecExpr
+    ParamSpecExpr, ARG_STAR2, Argument
 )
 from mypy.tvar_scope import TypeVarLikeScope
 from mypy.typevars import fill_typevars
@@ -91,7 +91,7 @@
     FunctionLike, UnboundType, TypeVarDef, TupleType, UnionType, StarType,
     CallableType, Overloaded, Instance, Type, AnyType, LiteralType, LiteralValue,
     TypeTranslator, TypeOfAny, TypeType, NoneType, PlaceholderType, TPDICT_NAMES, ProperType,
-    get_proper_type, get_proper_types, TypeAliasType)
+    get_proper_type, get_proper_types, TypeAliasType, TypedDictType)
 from mypy.typeops import function_type
 from mypy.type_visitor import TypeQuery
 from mypy.nodes import implicit_module_attrs
@@ -594,30 +594,50 @@ def analyze_func_def(self, defn: FuncDef) -> None:
                     defn.type = defn.type.copy_modified(ret_type=NoneType())
             self.prepare_method_signature(defn, self.type)
 
+        expand_kwargs = False
         # Analyze function signature
         with self.tvar_scope_frame(self.tvar_scope.method_frame()):
             if defn.type:
                 self.check_classvar_in_signature(defn.type)
                 assert isinstance(defn.type, CallableType)
+                expanded_kwargs_count = self.expanded_kwargs(defn)
+                if expanded_kwargs_count is not None:
+                    expand_kwargs = True
+
                 # Signature must be analyzed in the surrounding scope so that
                 # class-level imported names and type variables are in scope.
                 analyzer = self.type_analyzer()
                 tag = self.track_incomplete_refs()
-                result = analyzer.visit_callable_type(defn.type, nested=False)
+                result = analyzer.visit_callable_type(
+                    defn.type,
+                    nested=False,
+                    expand_kwargs=expand_kwargs)
                 # Don't store not ready types (including placeholders).
                 if self.found_incomplete_ref(tag) or has_placeholder(result):
                     self.defer(defn)
                     return
                 assert isinstance(result, ProperType)
                 defn.type = result
                 self.add_type_alias_deps(analyzer.aliases_used)
-                self.check_function_signature(defn)
+                self.check_function_signature(defn, expanded_kwargs=expanded_kwargs_count)
                 if isinstance(defn, FuncDef):
                     assert isinstance(defn.type, CallableType)
                     defn.type = set_callable_name(defn.type, defn)
 
         self.analyze_arg_initializers(defn)
-        self.analyze_function_body(defn)
+        self.analyze_function_body(defn, expand_kwargs=expand_kwargs)
+
+        if expand_kwargs:
+            # Modify defn's arguments.
+            assert isinstance(defn, FuncDef)
+            assert isinstance(defn.type, CallableType)
+            defn.arguments = [
+                Argument(Var(name, typ), typ, None, kind)
+                for name, typ, kind
+                in zip(defn.type.arg_names, defn.type.arg_types, defn.type.arg_kinds)
+                if name is not None
+            ]
+
         if (defn.is_coroutine and
                 isinstance(defn.type, CallableType) and
                 self.wrapped_coro_return_types.get(defn) != defn.type):
@@ -923,7 +943,7 @@ def analyze_arg_initializers(self, defn: FuncItem) -> None:
                 if arg.initializer:
                     arg.initializer.accept(self)
 
-    def analyze_function_body(self, defn: FuncItem) -> None:
+    def analyze_function_body(self, defn: FuncItem, expand_kwargs: bool = False) -> None:
         is_method = self.is_class_scope()
         with self.tvar_scope_frame(self.tvar_scope.method_frame()):
             # Bind the type variables again to visit the body.
@@ -933,6 +953,13 @@ def analyze_function_body(self, defn: FuncItem) -> None:
             self.function_stack.append(defn)
             self.enter(defn)
             for arg in defn.arguments:
+                if arg.kind == ARG_STAR2 and expand_kwargs:
+                    assert isinstance(arg.type_annotation, UnboundType)
+                    expand_type = arg.type_annotation.args[0]
+                    expand_arg = expand_type.accept(a)
+                    expand_var = Var(arg.variable.name, expand_arg)
+                    self.add_local(expand_var, defn)
+                    continue
                 self.add_local(arg.variable, defn)
 
             # The first argument of a non-static, non-class method is like 'self'
@@ -958,7 +985,9 @@ def check_classvar_in_signature(self, typ: ProperType) -> None:
                 # Show only one error per signature
                 break
 
-    def check_function_signature(self, fdef: FuncItem) -> None:
+    def check_function_signature(self,
+                                 fdef: FuncItem,
+                                 expanded_kwargs: Optional[int] = None) -> None:
         sig = fdef.type
         assert isinstance(sig, CallableType)
         if len(sig.arg_types) < len(fdef.arguments):
@@ -967,7 +996,12 @@ def check_function_signature(self, fdef: FuncItem) -> None:
             num_extra_anys = len(fdef.arguments) - len(sig.arg_types)
             extra_anys = [AnyType(TypeOfAny.from_error)] * num_extra_anys
             sig.arg_types.extend(extra_anys)
-        elif len(sig.arg_types) > len(fdef.arguments):
+        elif (
+            (expanded_kwargs is not None and
+            len(sig.arg_types) > len(fdef.arguments) + expanded_kwargs - 1) or
+            (expanded_kwargs is None and
+            len(sig.arg_types) > len(fdef.arguments))
+        ):
             self.fail('Type signature has too many arguments', fdef, blocker=True)
 
     def visit_decorator(self, dec: Decorator) -> None:
@@ -5030,6 +5064,43 @@ def set_future_import_flags(self, module_name: str) -> None:
     def is_future_flag_set(self, flag: str) -> bool:
         return flag in self.future_import_flags
 
+    def expanded_kwargs(self, defn: FuncDef) -> Optional[int]:
+        """Check if **kwargs are typed using the Expand special form.
+
+        If there are no **kwargs or they are typed without using Expand return None.
+        Otherwise return the number of items in the TypedDict.
+        """
+        assert isinstance(defn.type, CallableType)
+        if ARG_STAR2 in defn.type.arg_kinds:
+            kwargs_index = defn.type.arg_kinds.index(ARG_STAR2)
+            kwargs_type = defn.type.arg_types[kwargs_index]
+            kwargs_type = get_proper_type(kwargs_type)
+            if isinstance(kwargs_type, UnboundType):
+                sym = self.lookup_qualified(kwargs_type.name, kwargs_type)
+                if (sym is not None and
+                sym.node is not None and
+                sym.node.fullname == 'mypy_extensions.Expand'):
+                    if len(kwargs_type.args) != 1:
+                        self.fail(
+                            'Expand[...] must have exactly one type argument that is a TypedDict',
+                            kwargs_type
+                        )
+                        return None
+                    argument_type = get_proper_type(self.anal_type(kwargs_type.args[0]))
+                    if not isinstance(argument_type, TypedDictType):
+                        self.fail(
+                            'Expand[...] accepts only TypedDict as type argument',
+                            kwargs_type)
+                        return None
+                    if argument_type.items is None:
+                        self.fail(
+                            'Expand[...] cannot accept an empty TypedDict',
+                            kwargs_type
+                        )
+                        return None
+                    return len(argument_type.items)
+        return None
+
 
 class HasPlaceholders(TypeQuery[bool]):
     def __init__(self) -> None:

mypy/test/testcheck.py

@@ -95,6 +95,7 @@
     'check-generic-alias.test',
     'check-typeguard.test',
     'check-functools.test',
+    'check-expand.test',
 ]
 
 # Tests that use Python 3.8-only AST features (like expression-scoped ignores):

mypy/typeanal.py

@@ -348,6 +348,10 @@ 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 == 'mypy_extensions.Expand':
+            if self.nesting_level > 0:
+                self.fail("Invalid type: Expand[...] cannot be nested inside other type", t)
+            self.fail("Expand[...] can only be used with **kwargs in function declarations", t)
         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')
@@ -523,14 +527,19 @@ def visit_type_alias_type(self, t: TypeAliasType) -> Type:
     def visit_type_var(self, t: TypeVarType) -> Type:
         return t
 
-    def visit_callable_type(self, t: CallableType, nested: bool = True) -> Type:
+    def visit_callable_type(self,
+                            t: CallableType,
+                            nested: bool = True,
+                            expand_kwargs: bool = False) -> Type:
         # Every Callable can bind its own type variables, if they're not in the outer scope
         with self.tvar_scope_frame():
             if self.defining_alias:
                 variables = t.variables
             else:
                 variables = self.bind_function_type_variables(t, t)
             special = self.anal_type_guard(t.ret_type)
+            if expand_kwargs:
+                self.expand_kwargs(t)
             ret = t.copy_modified(arg_types=self.anal_array(t.arg_types, nested=nested),
                                   ret_type=self.anal_type(t.ret_type, nested=nested),
                                   # If the fallback isn't filled in yet,
@@ -542,6 +551,35 @@ def visit_callable_type(self, t: CallableType, nested: bool = True) -> Type:
                                   )
         return ret
 
+    def expand_kwargs(self, t: CallableType) -> None:
+        """Expand the function's signature with TypedDict items in place."""
+        kwargs_index = t.arg_kinds.index(ARG_STAR2)
+        expand = get_proper_type(t.arg_types[kwargs_index])
+        assert isinstance(expand, UnboundType)
+        # Get the TypedDict from Expand and analyze it.
+        kwargs_type = expand.args[0].accept(self)
+        kwargs_type = get_proper_type(kwargs_type)
+        assert isinstance(kwargs_type, TypedDictType)
+        assert kwargs_type.items is not None
+        # Expand the TypedDict.
+        expanded_types = []
+        expanded_kinds = []
+        expanded_names = []
+        for name, type in kwargs_type.items.items():
+            if name in t.arg_names:
+                self.fail(f'Cannot expand keyword arguments. Keyword "{name}" already present', t)
+            expanded_types.append(type)
+            expanded_kinds.append(
+                ARG_NAMED if name in kwargs_type.required_keys
+                else ARG_NAMED_OPT)
+            expanded_names.append(name)
+        del t.arg_types[kwargs_index]
+        del t.arg_kinds[kwargs_index]
+        del t.arg_names[kwargs_index]
+        t.arg_types += expanded_types
+        t.arg_kinds += expanded_kinds
+        t.arg_names += expanded_names
+
     def visit_type_guard_type(self, t: TypeGuardType) -> Type:
         return t
 

mypy/typeshed/stubs/mypy-extensions/mypy_extensions.pyi

@@ -1,6 +1,7 @@
 import abc
 import sys
 from typing import Any, Callable, Dict, Generic, ItemsView, KeysView, Mapping, Optional, Type, TypeVar, Union, ValuesView
+from typing import _SpecialForm  # type: ignore
 
 _T = TypeVar("_T")
 _U = TypeVar("_U")
@@ -44,3 +45,7 @@ def trait(cls: Any) -> Any: ...
 def mypyc_attr(*attrs: str, **kwattrs: object) -> Callable[[_T], _T]: ...
 
 class FlexibleAlias(Generic[_T, _U]): ...
+
+class _Expand(_SpecialForm): ... 
+
+Expand: _Expand = ...

setup.py

@@ -193,7 +193,7 @@ def run(self):
       # When changing this, also update mypy-requirements.txt.
       install_requires=["typed_ast >= 1.4.0, < 1.5.0; python_version<'3.8'",
                         'typing_extensions>=3.7.4',
-                        'mypy_extensions >= 0.4.3, < 0.5.0',
+                        'mypy_extensions >= 0.4.3, < 0.6.0',
                         'toml',
                         ],
       # Same here.

test-data/unit/check-expand.test

@@ -0,0 +1,122 @@
+-- Check Expand[...] for `**kwargs`.
+
+[case testExpandOutsideOfKwargs]
+from mypy_extensions import Expand, TypedDict
+class Person(TypedDict):
+    name: str
+    age: int
+
+x: Expand[Person]  # E: Expand[...] can only be used with **kwargs in function declarations # E: Variable "mypy_extensions.Expand" is not valid as a type # N: See https://mypy.readthedocs.io/en/stable/common_issues.html#variables-vs-type-aliases
+
+def foo(x: Expand[Person]) -> None:  # E: Expand[...] can only be used with **kwargs in function declarations # E: Variable "mypy_extensions.Expand" is not valid as a type # N: See https://mypy.readthedocs.io/en/stable/common_issues.html#variables-vs-type-aliases
+    ...
+
+def bar(x: int, *args: Expand[Person]) -> None:  # E: Expand[...] can only be used with **kwargs in function declarations # E: Variable "mypy_extensions.Expand" is not valid as a type # N: See https://mypy.readthedocs.io/en/stable/common_issues.html#variables-vs-type-aliases
+    ...
+
+# Should not result in error.
+def baz(**kwargs: Expand[Person]) -> None:
+    ...
+
+[builtins fixtures/dict.pyi]
+
+
+[case testExpandWithoutTypedDict]
+from mypy_extensions import Expand
+
+def foo(**kwargs: Expand[dict]) -> None:  # E: Expand[...] accepts only TypedDict as type argument # E: Expand[...] can only be used with **kwargs in function declarations # E: Variable "mypy_extensions.Expand" is not valid as a type # N: See https://mypy.readthedocs.io/en/stable/common_issues.html#variables-vs-type-aliases
+    ...
+
+[builtins fixtures/dict.pyi]
+
+
+[case testExpandsTypedDictTotality]
+from mypy_extensions import Expand, TypedDict
+
+class Circle(TypedDict, total=True):
+    radius: int
+    color: str
+    x: int
+    y: int
+
+def foo(**kwargs: Expand[Circle]):
+    ...
+
+foo(x=0, y=0, color='orange')  # E: Missing named argument "radius" for "foo"
+
+class Square(TypedDict, total=False):
+    side: int
+    color: str
+
+def bar(**kwargs: Expand[Square]):
+    ...
+
+# OK, because Square's totality is `False`.
+bar(side=12)
+
+[builtins fixtures/dict.pyi]
+
+
+[case testExpandUnexpectedKeyword]
+from mypy_extensions import Expand, TypedDict
+
+class Person(TypedDict, total=False):
+    name: str
+    age: int
+
+def foo(**kwargs: Expand[Person]) -> None:  # N: "foo" defined here
+    ...
+
+foo(name='John', age=42, department='Sales')  # E: Unexpected keyword argument "department" for "foo"
+
+foo(name='Jennifer', age=38)
+
+[builtins fixtures/dict.pyi]
+
+
+[case testExpandKeywordTypes]
+from mypy_extensions import Expand, TypedDict
+
+class Person(TypedDict):
+    name: str
+    age: int
+
+def foo(**kwargs: Expand[Person]):
+    ...
+
+foo(name='John', age='42')  # E: Argument "age" to "foo" has incompatible type "str"; expected "int"
+
+foo(name='Jennifer', age=38)
+
+[builtins fixtures/dict.pyi]
+
+
+[case testFunctionBodyWithExpandedKwargs]
+from mypy_extensions import Expand, TypedDict
+
+class Person(TypedDict):
+    name: str
+    age: int
+
+def foo(**kwargs: Expand[Person]) -> int:
+    name: str = kwargs['name']
+    age: str = kwargs['age']  # E: Incompatible types in assignment (expression has type "int", variable has type "str")
+    department: str = kwargs['department']  # E: TypedDict "Person" has no key "department"
+    return kwargs['age']
+
+foo(name='John', age=42)
+
+[builtins fixtures/dict.pyi]
+
+
+[case testExpandWithDuplicateKeywords]
+from mypy_extensions import Expand, TypedDict
+
+class Person(TypedDict):
+    name: str
+    age: int
+
+def foo(name: str, **kwargs: Expand[Person]) -> None:  # E: Cannot expand keyword arguments. Keyword "name" already present
+    ...
+
+[builtins fixtures/dict.pyi]