Type check typed dict as caller **kwargs

mypy/argmap.py

@@ -1,8 +1,8 @@
 """Utilities for mapping between actual and formal arguments (and their types)."""
 
-from typing import List, Optional, Sequence, Callable
+from typing import List, Optional, Sequence, Callable, Set
 
-from mypy.types import Type, Instance, TupleType, AnyType, TypeOfAny
+from mypy.types import Type, Instance, TupleType, AnyType, TypeOfAny, TypedDictType
 from mypy import nodes
 
 
@@ -65,13 +65,24 @@ def map_actuals_to_formals(caller_kinds: List[int],
                 map[callee_kinds.index(nodes.ARG_STAR2)].append(i)
         else:
             assert kind == nodes.ARG_STAR2
-            for j in range(ncallee):
-                # TODO tuple varargs complicate this
-                no_certain_match = (
-                    not map[j] or caller_kinds[map[j][0]] == nodes.ARG_STAR)
-                if ((callee_names[j] and no_certain_match)
-                        or callee_kinds[j] == nodes.ARG_STAR2):
-                    map[j].append(i)
+            argt = caller_arg_type(i)
+            if isinstance(argt, TypedDictType):
+                for name, value in argt.items.items():
+                    if name in callee_names:
+                        map[callee_names.index(name)].append(i)
+                    elif nodes.ARG_STAR2 in callee_kinds:
+                        map[callee_kinds.index(nodes.ARG_STAR2)].append(i)
+            else:
+                # We don't exactly know which **kwargs are provided by the
+                # caller. Assume that they will fill the remaining arguments.
+                for j in range(ncallee):
+                    # TODO: If there are also tuple varargs, we might be missing some potential
+                    #       matches if the tuple was short enough to not match everything.
+                    no_certain_match = (
+                        not map[j] or caller_kinds[map[j][0]] == nodes.ARG_STAR)
+                    if ((callee_names[j] and no_certain_match)
+                            or callee_kinds[j] == nodes.ARG_STAR2):
+                        map[j].append(i)
     return map
 
 
@@ -95,35 +106,87 @@ def map_formals_to_actuals(caller_kinds: List[int],
     return actual_to_formal
 
 
-def get_actual_type(arg_type: Type, kind: int,
-                    tuple_counter: List[int]) -> Type:
-    """Return the type of an actual argument with the given kind.
+class ArgTypeExpander:
+    """Utility class for mapping actual argument types to formal arguments.
+
+    One of the main responsibilities is to expand caller tuple *args and TypedDict
+    **kwargs, and to keep track of which tuple/TypedDict items have already been
+    consumed.
+
+    Example:
+
+       def f(x: int, *args: str) -> None: ...
+       f(*(1, 'x', 1.1))
+
+    We'd call expand_actual_type three times:
 
-    If the argument is a *arg, return the individual argument item.
+      1. The first call would provide 'int' as the actual type of 'x' (from '1').
+      2. The second call would provide 'str' as one of the actual types for '*args'.
+      2. The third call would provide 'float' as one of the actual types for '*args'.
+
+    A single instance can process all the arguments for a single call. Each call
+    needs a separate instance since instances have per-call state.
     """
 
-    if kind == nodes.ARG_STAR:
-        if isinstance(arg_type, Instance):
-            if arg_type.type.fullname() == 'builtins.list':
-                # List *arg.
-                return arg_type.args[0]
-            elif arg_type.args:
-                # TODO try to map type arguments to Iterable
-                return arg_type.args[0]
+    def __init__(self) -> None:
+        # Next tuple *args index to use.
+        self.tuple_index = 0
+        # Keyword arguments in TypedDict **kwargs used.
+        self.kwargs_used = set()  # type: Set[str]
+
+    def expand_actual_type(self,
+                           actual_type: Type,
+                           actual_kind: int,
+                           formal_name: Optional[str],
+                           formal_kind: int) -> Type:
+        """Return the actual (caller) type(s) of a formal argument with the given kinds.
+
+        If the actual argument is a tuple *args, return the next individual tuple item that
+        maps to the formal arg.
+
+        If the actual argument is a TypedDict **kwargs, return the next matching typed dict
+        value type based on formal argument name and kind.
+
+        This is supposed to be called for each formal, in order. Call multiple times per
+        formal if multiple actuals map to a formal.
+        """
+        if actual_kind == nodes.ARG_STAR:
+            if isinstance(actual_type, Instance):
+                if actual_type.type.fullname() == 'builtins.list':
+                    # List *arg.
+                    return actual_type.args[0]
+                elif actual_type.args:
+                    # TODO: Try to map type arguments to Iterable
+                    return actual_type.args[0]
+                else:
+                    return AnyType(TypeOfAny.from_error)
+            elif isinstance(actual_type, TupleType):
+                # Get the next tuple item of a tuple *arg.
+                if self.tuple_index >= len(actual_type.items):
+                    # Exhausted a tuple -- continue to the next *args.
+                    self.tuple_index = 1
+                else:
+                    self.tuple_index += 1
+                return actual_type.items[self.tuple_index - 1]
+            else:
+                return AnyType(TypeOfAny.from_error)
+        elif actual_kind == nodes.ARG_STAR2:
+            if isinstance(actual_type, TypedDictType):
+                if formal_kind != nodes.ARG_STAR2 and formal_name in actual_type.items:
+                    # Lookup type based on keyword argument name.
+                    assert formal_name is not None
+                else:
+                    # Pick an arbitrary item if no specified keyword is expected.
+                    formal_name = (set(actual_type.items.keys()) - self.kwargs_used).pop()
+                self.kwargs_used.add(formal_name)
+                return actual_type.items[formal_name]
+            elif (isinstance(actual_type, Instance)
+                  and (actual_type.type.fullname() == 'builtins.dict')):
+                # Dict **arg.
+                # TODO: Handle arbitrary Mapping
+                return actual_type.args[1]
             else:
                 return AnyType(TypeOfAny.from_error)
-        elif isinstance(arg_type, TupleType):
-            # Get the next tuple item of a tuple *arg.
-            tuple_counter[0] += 1
-            return arg_type.items[tuple_counter[0] - 1]
-        else:
-            return AnyType(TypeOfAny.from_error)
-    elif kind == nodes.ARG_STAR2:
-        if isinstance(arg_type, Instance) and (arg_type.type.fullname() == 'builtins.dict'):
-            # Dict **arg. TODO more general (Mapping)
-            return arg_type.args[1]
         else:
-            return AnyType(TypeOfAny.from_error)
-    else:
-        # No translation for other kinds.
-        return arg_type
+            # No translation for other kinds -- 1:1 mapping.
+            return actual_type