Basic support for generic TypedDicts

ilevkivskyi · Aug 15, 2022 · Aug 11, 2022 · Aug 11, 2022 · Aug 11, 2022 · Aug 11, 2022
commit 7d3220eda34113d03a12d03b3c29b05e0a94201e
diff --git a/mypy/checkexpr.py b/mypy/checkexpr.py
@@ -318,7 +318,14 @@ def analyze_ref_expr(self, e: RefExpr, lvalue: bool = False) -> Type:
             result = node.type
         elif isinstance(node, TypeInfo):
             # Reference to a type object.
-            result = type_object_type(node, self.named_type)
+            if node.typeddict_type:
+                # Use alias target to get the correct fallback (not just typing.TypedDict).
+                assert node.special_alias is not None
+>
+                assert isinstance(target, ProperType) and isinstance(target, TypedDictType)
+                result = self.typeddict_callable(target)
+            else:
+                result = type_object_type(node, self.named_type)
             if isinstance(result, CallableType) and isinstance(  # type: ignore
                 result.ret_type, Instance
             ):
@@ -396,6 +403,20 @@ def visit_call_expr_inner(self, e: CallExpr, allow_none_return: bool = False) ->
                 fallback=Instance(e.callee.node, [])
             )
             return self.check_typeddict_call(typeddict_type, e.arg_kinds, e.arg_names, e.args, e)
+        if (
+            isinstance(e.callee, IndexExpr)
+            and isinstance(e.callee.base, RefExpr)
+            and isinstance(e.callee.base.node, TypeInfo)
+            and e.callee.base.node.typeddict_type is not None
+        ):
+            # Apply type argument form type application.
+            typeddict_callable = get_proper_type(self.accept(e.callee))
+            if isinstance(typeddict_callable, CallableType):
+                typeddict_applied = get_proper_type(typeddict_callable.ret_type)
+                assert isinstance(typeddict_applied, TypedDictType)
+                return self.check_typeddict_call(
+                    typeddict_applied, e.arg_kinds, e.arg_names, e.args, e
+                )
         if (
             isinstance(e.callee, NameExpr)
             and e.callee.name in ("isinstance", "issubclass")
@@ -692,6 +713,24 @@ def check_typeddict_call_with_dict(
         else:
             return AnyType(TypeOfAny.from_error)
 
+    def typeddict_callable(self, typed_dict: TypedDictType) -> CallableType:
+        """Contsruct a reasonable type for a TypedDict type in runtime context.
+
+        If it appears as a callee, it will be special-cased anyway, e.g. it is
+        also allowed to accept a single positional argument if it is a dict literal.
+        """
+        expected_types = list(typed_dict.items.values())
+        kinds = [ArgKind.ARG_NAMED] * len(expected_types)
+        names = list(typed_dict.items.keys())
+        return CallableType(
+            expected_types,
+            kinds,
+            names,
+            typed_dict,
+            self.named_type("builtins.type"),
+            variables=typed_dict.fallback.type.defn.type_vars,
+        )
+
     def check_typeddict_call_with_kwargs(
         self, callee: TypedDictType, kwargs: Dict[str, Expression], context: Context
     ) -> Type:
@@ -707,7 +746,27 @@ def check_typeddict_call_with_kwargs(
             )
             return AnyType(TypeOfAny.from_error)
 
-        for (item_name, item_expected_type) in callee.items.items():
+        # We don't show any errors, just infer types in a generic TypedDict type,
+        # a custom error message will be given below, if there are errors.
+        with self.msg.filter_errors(), self.chk.local_type_map():
+            orig_ret_type, _ = self.check_callable_call(
+                self.typeddict_callable(callee),
+                list(kwargs.values()),
+                [ArgKind.ARG_NAMED] * len(kwargs),
+                context,
+                list(kwargs.keys()),
+                None,
+                None,
+                None,
+            )
+
+        ret_type = get_proper_type(orig_ret_type)
+        if not isinstance(ret_type, TypedDictType):
+            # If something went really wrong, type-check call with original type,
+            # this may give a better error message.
+            ret_type = callee
+
+        for (item_name, item_expected_type) in ret_type.items.items():
             if item_name in kwargs:
                 item_value = kwargs[item_name]
                 self.chk.check_simple_assignment(
@@ -720,7 +779,7 @@ def check_typeddict_call_with_kwargs(
                     code=codes.TYPEDDICT_ITEM,
                 )
 
-        return callee
+        return orig_ret_type
 
     def get_partial_self_var(self, expr: MemberExpr) -> Optional[Var]:
         """Get variable node for a partial self attribute.

diff --git a/mypy/semanal.py b/mypy/semanal.py
@@ -1379,6 +1379,12 @@ def analyze_class(self, defn: ClassDef) -> None:
             return
 
         if self.analyze_typeddict_classdef(defn):
+            if defn.info:
+                defn.type_vars = tvar_defs
+                defn.info.type_vars = []
+                defn.info.add_type_vars()
+                assert defn.info.special_alias is not None
+                defn.info.special_alias.alias_tvars = list(defn.info.type_vars)
             return
 
         if self.analyze_namedtuple_classdef(defn):

diff --git a/mypy/typeanal.py b/mypy/typeanal.py
@@ -621,12 +621,9 @@ def analyze_type_with_type_info(
         if td is not None:
             # The class has a TypedDict[...] base class so it will be
             # represented as a typeddict type.
-            if args:
-                self.fail("Generic TypedDict types not supported", ctx)
-                return AnyType(TypeOfAny.from_error)
             if info.special_alias:
                 # We don't support generic TypedDict types yet.
-                return TypeAliasType(info.special_alias, [])
+                return TypeAliasType(info.special_alias, self.anal_array(args))
             # Create a named TypedDictType
             return td.copy_modified(
                 item_types=self.anal_array(list(td.items.values())), fallback=instance

diff --git a/mypy/types.py b/mypy/types.py
@@ -1730,6 +1730,8 @@ def type_object(self) -> mypy.nodes.TypeInfo:
             ret = get_proper_type(ret.upper_bound)
         if isinstance(ret, TupleType):
             ret = ret.partial_fallback
+        if isinstance(ret, TypedDictType):
+            ret = ret.fallback
         assert isinstance(ret, Instance)
         return ret.type
 

diff --git a/test-data/unit/check-serialize.test b/test-data/unit/check-serialize.test
@@ -1066,11 +1066,11 @@ class C:
 [out1]
 main:2: note: Revealed type is "TypedDict('ntcrash.C.A@4', {'x': builtins.int})"
 main:3: note: Revealed type is "TypedDict('ntcrash.C.A@4', {'x': builtins.int})"
-main:4: note: Revealed type is "def () -> ntcrash.C.A@4"
+main:4: note: Revealed type is "def (*, x: builtins.int) -> TypedDict('ntcrash.C.A@4', {'x': builtins.int})"
 [out2]
 main:2: note: Revealed type is "TypedDict('ntcrash.C.A@4', {'x': builtins.int})"
 main:3: note: Revealed type is "TypedDict('ntcrash.C.A@4', {'x': builtins.int})"
-main:4: note: Revealed type is "def () -> ntcrash.C.A@4"
+main:4: note: Revealed type is "def (*, x: builtins.int) -> TypedDict('ntcrash.C.A@4', {'x': builtins.int})"
 
 [case testSerializeNonTotalTypedDict]
 from m import d

diff --git a/test-data/unit/check-typeddict.test b/test-data/unit/check-typeddict.test
@@ -789,7 +789,7 @@ from mypy_extensions import TypedDict
 D = TypedDict('D', {'x': int})
 d: object
 if isinstance(d, D):   # E: Cannot use isinstance() with TypedDict type
-    reveal_type(d)     # N: Revealed type is "__main__.D"
+    reveal_type(d)     # N: Revealed type is "TypedDict('__main__.D', {'x': builtins.int})"
 issubclass(object, D)  # E: Cannot use issubclass() with TypedDict type
 [builtins fixtures/isinstancelist.pyi]
 
@@ -1517,7 +1517,7 @@ from b import tp
 x: tp
 reveal_type(x['x'])  # N: Revealed type is "builtins.int"
 
-reveal_type(tp)  # N: Revealed type is "def () -> b.tp"
+reveal_type(tp)  # N: Revealed type is "def (*, x: builtins.int) -> TypedDict('b.tp', {'x': builtins.int})"
 tp(x='no')  # E: Incompatible types (expression has type "str", TypedDict item "x" has type "int")
 
 [file b.py]
@@ -2395,3 +2395,21 @@ def func(foo: Union[F1, F2]):
                   # E: Argument 1 to "__setitem__" has incompatible type "int"; expected "str"
 [builtins fixtures/dict.pyi]
 [typing fixtures/typing-typeddict.pyi]
+
+[case testGenericTypedDict]
+from typing import TypedDict, Generic, TypeVar
+
+T = TypeVar("T")
+
+class TD(TypedDict, Generic[T]):
+    key: int
+    value: T
+
+tds: TD[str]
+reveal_type(tds)  # N: Revealed type is "TypedDict('__main__.TD', {'key': builtins.int, 'value': builtins.str})"
+
+tdi = TD(key=0, value=0)
+reveal_type(tdi)  # N: Revealed type is "TypedDict('__main__.TD', {'key': builtins.int, 'value': builtins.int})"
+TD[str](key=0, value=0)  # E: Incompatible types (expression has type "int", TypedDict item "value" has type "str")
+[builtins fixtures/dict.pyi]
+[typing fixtures/typing-typeddict.pyi]