Support TypedDicts with missing keys (total=False) (#3558)

* Basic support for TypedDicts with missing keys (total=False)

Only the functional syntax is supported.

* Support get(key, {}) and fix construction of partial typed dict

* Fix subtyping of non-total typed dicts

* Fix join with non-total typed dict

* Fix meet with non-total typed dicts

* Add serialization test case

* Support TypedDict total keyword argument with class syntax

* Attempt to fix Python 3.3

* Add minimal runtime `total` support to mypy_extensions

There is no support for introspection of `total` yet.

* Fix tests on pre-3.6 Python and improve introspection

Make TypedDict `total` introspectable.

* Fix lint

* Fix problems caused by merge

* Allow td['key'] even if td is not total

* Fix lint

* Add test case

* Address review feedback

* Update comment

Author: JukkaL

extensions/mypy_extensions.py

@@ -30,16 +30,18 @@ def _dict_new(cls, *args, **kwargs):
 
 
 def _typeddict_new(cls, _typename, _fields=None, **kwargs):
+    total = kwargs.pop('total', True)
     if _fields is None:
         _fields = kwargs
     elif kwargs:
         raise TypeError("TypedDict takes either a dict or keyword arguments,"
                         " but not both")
-    return _TypedDictMeta(_typename, (), {'__annotations__': dict(_fields)})
+    return _TypedDictMeta(_typename, (), {'__annotations__': dict(_fields),
+                                          '__total__': total})
 
 
 class _TypedDictMeta(type):
-    def __new__(cls, name, bases, ns):
+    def __new__(cls, name, bases, ns, total=True):
         # Create new typed dict class object.
         # This method is called directly when TypedDict is subclassed,
         # or via _typeddict_new when TypedDict is instantiated. This way
@@ -59,6 +61,8 @@ def __new__(cls, name, bases, ns):
         for base in bases:
             anns.update(base.__dict__.get('__annotations__', {}))
         tp_dict.__annotations__ = anns
+        if not hasattr(tp_dict, '__total__'):
+            tp_dict.__total__ = total
         return tp_dict
 
     __instancecheck__ = __subclasscheck__ = _check_fails

mypy/checkexpr.py

@@ -292,31 +292,31 @@ def check_typeddict_call_with_dict(self, callee: TypedDictType,
     def check_typeddict_call_with_kwargs(self, callee: TypedDictType,
                                          kwargs: 'OrderedDict[str, Expression]',
                                          context: Context) -> Type:
-        if callee.items.keys() != kwargs.keys():
-            callee_item_names = callee.items.keys()
-            kwargs_item_names = kwargs.keys()
-
+        if not (callee.required_keys <= set(kwargs.keys()) <= set(callee.items.keys())):
+            expected_item_names = [key for key in callee.items.keys()
+                                   if key in callee.required_keys or key in kwargs.keys()]
+            actual_item_names = kwargs.keys()
             self.msg.typeddict_instantiated_with_unexpected_items(
-                expected_item_names=list(callee_item_names),
-                actual_item_names=list(kwargs_item_names),
+                expected_item_names=list(expected_item_names),
+                actual_item_names=list(actual_item_names),
                 context=context)
             return AnyType()
 
         items = OrderedDict()  # type: OrderedDict[str, Type]
         for (item_name, item_expected_type) in callee.items.items():
-            item_value = kwargs[item_name]
-
-            self.chk.check_simple_assignment(
-                lvalue_type=item_expected_type, rvalue=item_value, context=item_value,
-                msg=messages.INCOMPATIBLE_TYPES,
-                lvalue_name='TypedDict item "{}"'.format(item_name),
-                rvalue_name='expression')
+            if item_name in kwargs:
+                item_value = kwargs[item_name]
+                self.chk.check_simple_assignment(
+                    lvalue_type=item_expected_type, rvalue=item_value, context=item_value,
+                    msg=messages.INCOMPATIBLE_TYPES,
+                    lvalue_name='TypedDict item "{}"'.format(item_name),
+                    rvalue_name='expression')
             items[item_name] = item_expected_type
 
         mapping_value_type = join.join_type_list(list(items.values()))
         fallback = self.chk.named_generic_type('typing.Mapping',
                                                [self.chk.str_type(), mapping_value_type])
-        return TypedDictType(items, fallback)
+        return TypedDictType(items, set(callee.required_keys), fallback)
 
     # Types and methods that can be used to infer partial types.
     item_args = {'builtins.list': ['append'],

mypy/fastparse.py

@@ -467,12 +467,15 @@ def visit_ClassDef(self, n: ast3.ClassDef) -> ClassDef:
             metaclass = stringify_name(metaclass_arg.value)
             if metaclass is None:
                 metaclass = '<error>'  # To be reported later
+        keywords = [(kw.arg, self.visit(kw.value))
+                    for kw in n.keywords]
 
         cdef = ClassDef(n.name,
                         self.as_block(n.body, n.lineno),
                         None,
                         self.translate_expr_list(n.bases),
-                        metaclass=metaclass)
+                        metaclass=metaclass,
+                        keywords=keywords)
         cdef.decorators = self.translate_expr_list(n.decorator_list)
         self.class_nesting -= 1
         return cdef

mypy/join.py

@@ -228,11 +228,15 @@ def visit_typeddict_type(self, t: TypedDictType) -> Type:
             items = OrderedDict([
                 (item_name, s_item_type)
                 for (item_name, s_item_type, t_item_type) in self.s.zip(t)
-                if is_equivalent(s_item_type, t_item_type)
+                if (is_equivalent(s_item_type, t_item_type) and
+                    (item_name in t.required_keys) == (item_name in self.s.required_keys))
             ])
             mapping_value_type = join_type_list(list(items.values()))
             fallback = self.s.create_anonymous_fallback(value_type=mapping_value_type)
-            return TypedDictType(items, fallback)
+            # We need to filter by items.keys() since some required keys present in both t and
+            # self.s might be missing from the join if the types are incompatible.
+            required_keys = set(items.keys()) & t.required_keys & self.s.required_keys
+            return TypedDictType(items, required_keys, fallback)
         elif isinstance(self.s, Instance):
             return join_instances(self.s, t.fallback)
         else:

mypy/meet.py

@@ -252,8 +252,9 @@ def visit_tuple_type(self, t: TupleType) -> Type:
 
     def visit_typeddict_type(self, t: TypedDictType) -> Type:
         if isinstance(self.s, TypedDictType):
-            for (_, l, r) in self.s.zip(t):
-                if not is_equivalent(l, r):
+            for (name, l, r) in self.s.zip(t):
+                if (not is_equivalent(l, r) or
+                        (name in t.required_keys) != (name in self.s.required_keys)):
                     return self.default(self.s)
             item_list = []  # type: List[Tuple[str, Type]]
             for (item_name, s_item_type, t_item_type) in self.s.zipall(t):
@@ -266,7 +267,8 @@ def visit_typeddict_type(self, t: TypedDictType) -> Type:
             items = OrderedDict(item_list)
             mapping_value_type = join_type_list(list(items.values()))
             fallback = self.s.create_anonymous_fallback(value_type=mapping_value_type)
-            return TypedDictType(items, fallback)
+            required_keys = t.required_keys | self.s.required_keys
+            return TypedDictType(items, required_keys, fallback)
         else:
             return self.default(self.s)
 

mypy/nodes.py

@@ -2,6 +2,7 @@
 
 import os
 from abc import abstractmethod
+from collections import OrderedDict
 
 from typing import (
     Any, TypeVar, List, Tuple, cast, Set, Dict, Union, Optional, Callable,
@@ -730,6 +731,7 @@ class ClassDef(Statement):
     info = None  # type: TypeInfo  # Related TypeInfo
     metaclass = ''  # type: Optional[str]
     decorators = None  # type: List[Expression]
+    keywords = None  # type: OrderedDict[str, Expression]
     analyzed = None  # type: Optional[Expression]
     has_incompatible_baseclass = False
 
@@ -738,13 +740,15 @@ def __init__(self,
                  defs: 'Block',
                  type_vars: List['mypy.types.TypeVarDef'] = None,
                  base_type_exprs: List[Expression] = None,
-                 metaclass: str = None) -> None:
+                 metaclass: str = None,
+                 keywords: List[Tuple[str, Expression]] = None) -> None:
         self.name = name
         self.defs = defs
         self.type_vars = type_vars or []
         self.base_type_exprs = base_type_exprs or []
         self.metaclass = metaclass
         self.decorators = []
+        self.keywords = OrderedDict(keywords or [])
 
     def accept(self, visitor: StatementVisitor[T]) -> T:
         return visitor.visit_class_def(self)

mypy/plugin.py

@@ -1,9 +1,10 @@
 """Plugin system for extending mypy."""
 
+from collections import OrderedDict
 from abc import abstractmethod
 from typing import Callable, List, Tuple, Optional, NamedTuple, TypeVar
 
-from mypy.nodes import Expression, StrExpr, IntExpr, UnaryExpr, Context
+from mypy.nodes import Expression, StrExpr, IntExpr, UnaryExpr, Context, DictExpr
 from mypy.types import (
     Type, Instance, CallableType, TypedDictType, UnionType, NoneTyp, FunctionLike, TypeVarType,
     AnyType, TypeList, UnboundType
@@ -263,17 +264,26 @@ def typed_dict_get_signature_callback(ctx: MethodSigContext) -> CallableType:
             and len(ctx.args[0]) == 1
             and isinstance(ctx.args[0][0], StrExpr)
             and len(signature.arg_types) == 2
-            and len(signature.variables) == 1):
+            and len(signature.variables) == 1
+            and len(ctx.args[1]) == 1):
         key = ctx.args[0][0].value
         value_type = ctx.type.items.get(key)
+        ret_type = signature.ret_type
         if value_type:
+            default_arg = ctx.args[1][0]
+            if (isinstance(value_type, TypedDictType)
+                    and isinstance(default_arg, DictExpr)
+                    and len(default_arg.items) == 0):
+                # Caller has empty dict {} as default for typed dict.
+                value_type = value_type.copy_modified(required_keys=set())
             # Tweak the signature to include the value type as context. It's
             # only needed for type inference since there's a union with a type
             # variable that accepts everything.
             tv = TypeVarType(signature.variables[0])
             return signature.copy_modified(
                 arg_types=[signature.arg_types[0],
-                           UnionType.make_simplified_union([value_type, tv])])
+                           UnionType.make_simplified_union([value_type, tv])],
+                ret_type=ret_type)
     return signature
 
 
@@ -288,8 +298,15 @@ def typed_dict_get_callback(ctx: MethodContext) -> Type:
             if value_type:
                 if len(ctx.arg_types) == 1:
                     return UnionType.make_simplified_union([value_type, NoneTyp()])
-                elif len(ctx.arg_types) == 2 and len(ctx.arg_types[1]) == 1:
-                    return UnionType.make_simplified_union([value_type, ctx.arg_types[1][0]])
+                elif (len(ctx.arg_types) == 2 and len(ctx.arg_types[1]) == 1
+                      and len(ctx.args[1]) == 1):
+                    default_arg = ctx.args[1][0]
+                    if (isinstance(default_arg, DictExpr) and len(default_arg.items) == 0
+                            and isinstance(value_type, TypedDictType)):
+                        # Special case '{}' as the default for a typed dict type.
+                        return value_type.copy_modified(required_keys=set())
+                    else:
+                        return UnionType.make_simplified_union([value_type, ctx.arg_types[1][0]])
             else:
                 ctx.api.msg.typeddict_item_name_not_found(ctx.type, key, ctx.context)
                 return AnyType()