Add support for call syntax

ilevkivskyi · Aug 15, 2022 · Aug 11, 2022 · Aug 11, 2022 · Aug 11, 2022 · Aug 11, 2022
commit cf0ba973db92f99f610f70cac6447cc9fb2c1a89
diff --git a/mypy/checkexpr.py b/mypy/checkexpr.py
@@ -730,6 +730,9 @@ def typeddict_callable(self, info: TypeInfo) -> CallableType:
 
         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.
+
+        Note it is not safe to move this to type_object_type() since it will crash
+        on plugin-generated TypedDicts, that may not have the special_alias.
         """
         assert info.special_alias is not None
 >
@@ -782,6 +785,7 @@ def check_typeddict_call_with_kwargs(
             if callee.fallback.type.special_alias is not None:
                 infer_callee = self.typeddict_callable(callee.fallback.type)
             else:
+                # Likely a TypedDict type generated by a plugin.
                 infer_callee = self.typeddict_callable_from_context(callee)
 
         # We don't show any errors, just infer types in a generic TypedDict type,

diff --git a/mypy/semanal.py b/mypy/semanal.py
@@ -1381,14 +1381,8 @@ def analyze_class(self, defn: ClassDef) -> None:
 
         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)
->
-                assert isinstance(target, ProperType) and isinstance(target, TypedDictType)
-                target.fallback.args = tuple(defn.type_vars)
+                self.setup_type_vars(defn, tvar_defs)
+                self.setup_alias_type_vars(defn)
             return
 
         if self.analyze_namedtuple_classdef(defn):
@@ -1413,6 +1407,19 @@ def analyze_class(self, defn: ClassDef) -> None:
                 self.analyze_class_decorator(defn, decorator)
             self.analyze_class_body_common(defn)
 
+    def setup_type_vars(self, defn: ClassDef, tvar_defs: List[TypeVarLikeType]) -> None:
+        defn.type_vars = tvar_defs
+        defn.info.type_vars = []
+        # we want to make sure any additional logic in add_type_vars gets run
+        defn.info.add_type_vars()
+
+    def setup_alias_type_vars(self, defn: ClassDef) -> None:
+        assert defn.info.special_alias is not None
+        defn.info.special_alias.alias_tvars = list(defn.info.type_vars)
+>
+        assert isinstance(target, ProperType) and isinstance(target, TypedDictType)
+        target.fallback.args = tuple(defn.type_vars)
+
     def is_core_builtin_class(self, defn: ClassDef) -> bool:
         return self.cur_mod_id == "builtins" and defn.name in CORE_BUILTIN_CLASSES
 
@@ -1689,6 +1696,29 @@ def get_all_bases_tvars(
                 tvars.extend(base_tvars)
         return remove_dups(tvars)
 
+    def get_and_bind_all_tvars(self, type_exprs: List[Expression]) -> List[TypeVarLikeType]:
+        """Return all type variable references in item type expressions.
+
+        This is a helper for generic TypedDicts and NamedTuples. Essentially it is
+        a simplified version of the logic we use for ClassDef bases. We duplicate
+        some amount of code, because it is hard to refactor common pieces.
+        """
+        tvars = []
+        for base_expr in type_exprs:
+            try:
+                base = self.expr_to_unanalyzed_type(base_expr)
+            except TypeTranslationError:
+                # This error will be caught later.
+                continue
+            base_tvars = base.accept(TypeVarLikeQuery(self.lookup_qualified, self.tvar_scope))
+            tvars.extend(base_tvars)
+        tvars = remove_dups(tvars)  # Variables are defined in order of textual appearance.
+        tvar_defs = []
+        for name, tvar_expr in tvars:
+            tvar_def = self.tvar_scope.bind_new(name, tvar_expr)
+            tvar_defs.append(tvar_def)
+        return tvar_defs
+
     def prepare_class_def(self, defn: ClassDef, info: Optional[TypeInfo] = None) -> None:
         """Prepare for the analysis of a class definition.
 
@@ -2693,7 +2723,7 @@ def analyze_typeddict_assign(self, s: AssignmentStmt) -> bool:
             return False
         lvalue = s.lvalues[0]
         name = lvalue.name
-        is_typed_dict, info = self.typed_dict_analyzer.check_typeddict(
+        is_typed_dict, info, tvar_defs = self.typed_dict_analyzer.check_typeddict(
             s.rvalue, name, self.is_func_scope()
         )
         if not is_typed_dict:
@@ -2704,6 +2734,10 @@ def analyze_typeddict_assign(self, s: AssignmentStmt) -> bool:
         # Yes, it's a valid typed dict, but defer if it is not ready.
         if not info:
             self.mark_incomplete(name, lvalue, becomes_typeinfo=True)
+        else:
+            defn = info.defn
+            self.setup_type_vars(defn, tvar_defs)
+            self.setup_alias_type_vars(defn)
         return True
 
     def analyze_lvalues(self, s: AssignmentStmt) -> None:

diff --git a/mypy/semanal_shared.py b/mypy/semanal_shared.py
@@ -32,6 +32,7 @@
     TupleType,
     Type,
     TypeVarId,
+    TypeVarLikeType,
     get_proper_type,
 )
 
@@ -158,6 +159,10 @@ def anal_type(
     ) -> Optional[Type]:
         raise NotImplementedError
 
+    @abstractmethod
+    def get_and_bind_all_tvars(self, type_exprs: List[Expression]) -> List[TypeVarLikeType]:
+        raise NotImplementedError
+
     @abstractmethod
     def basic_new_typeinfo(self, name: str, basetype_or_fallback: Instance, line: int) -> TypeInfo:
         raise NotImplementedError

diff --git a/mypy/semanal_typeddict.py b/mypy/semanal_typeddict.py
@@ -38,6 +38,7 @@
     Type,
     TypedDictType,
     TypeOfAny,
+    TypeVarLikeType,
     replace_alias_tvars,
 )
 
@@ -316,7 +317,7 @@ def analyze_typeddict_classdef_fields(
 
     def check_typeddict(
         self, node: Expression, var_name: Optional[str], is_func_scope: bool
-    ) -> Tuple[bool, Optional[TypeInfo]]:
+    ) -> Tuple[bool, Optional[TypeInfo], List[TypeVarLikeType]]:
         """Check if a call defines a TypedDict.
 
         The optional var_name argument is the name of the variable to
@@ -329,20 +330,20 @@ def check_typeddict(
         return (True, None).
         """
         if not isinstance(node, CallExpr):
-            return False, None
+            return False, None, []
         call = node
         callee = call.callee
         if not isinstance(callee, RefExpr):
-            return False, None
+            return False, None, []
         fullname = callee.fullname
         if fullname not in TPDICT_NAMES:
-            return False, None
+            return False, None, []
         res = self.parse_typeddict_args(call)
         if res is None:
             # This is a valid typed dict, but some type is not ready.
             # The caller should defer this until next iteration.
-            return True, None
-        name, items, types, total, ok = res
+            return True, None, []
+        name, items, types, total, tvar_defs, ok = res
         if not ok:
             # Error. Construct dummy return value.
             info = self.build_typeddict_typeinfo("TypedDict", [], [], set(), call.line, None)
@@ -381,11 +382,11 @@ def check_typeddict(
             self.api.add_symbol(var_name, info, node)
         call.analyzed = TypedDictExpr(info)
         call.analyzed.set_line(call)
-        return True, info
+        return True, info, tvar_defs
 
     def parse_typeddict_args(
         self, call: CallExpr
-    ) -> Optional[Tuple[str, List[str], List[Type], bool, bool]]:
+    ) -> Optional[Tuple[str, List[str], List[Type], bool, List[TypeVarLikeType], bool]]:
         """Parse typed dict call expression.
 
         Return names, types, totality, was there an error during parsing.
@@ -420,6 +421,7 @@ def parse_typeddict_args(
                     'TypedDict() "total" argument must be True or False', call
                 )
         dictexpr = args[1]
+        tvar_defs = self.api.get_and_bind_all_tvars([t for k, t in dictexpr.items])
         res = self.parse_typeddict_fields_with_types(dictexpr.items, call)
         if res is None:
             # One of the types is not ready, defer.
@@ -435,7 +437,7 @@ def parse_typeddict_args(
                 if has_any_from_unimported_type(t):
                     self.msg.unimported_type_becomes_any("Type of a TypedDict key", t, dictexpr)
         assert total is not None
-        return args[0].value, items, types, total, ok
+        return args[0].value, items, types, total, tvar_defs, ok
 
     def parse_typeddict_fields_with_types(
         self, dict_items: List[Tuple[Optional[Expression], Expression]], context: Context
@@ -488,9 +490,9 @@ def parse_typeddict_fields_with_types(
 
     def fail_typeddict_arg(
         self, message: str, context: Context
-    ) -> Tuple[str, List[str], List[Type], bool, bool]:
+    ) -> Tuple[str, List[str], List[Type], bool, List[TypeVarLikeType], bool]:
         self.fail(message, context)
-        return "", [], [], True, False
+        return "", [], [], True, [], False
 
     def build_typeddict_typeinfo(
         self,

diff --git a/test-data/unit/check-typeddict.test b/test-data/unit/check-typeddict.test
@@ -2507,3 +2507,20 @@ Alias(key=0, value=0)  # E: Missing type parameters for generic type "Alias" \
                        # E: Incompatible types (expression has type "int", TypedDict item "value" has type "List[Any]")
 [builtins fixtures/dict.pyi]
 [typing fixtures/typing-typeddict.pyi]
+
+[case testGenericTypedDictCallSyntax]
+from typing import TypedDict, TypeVar
+
+T = TypeVar("T")
+TD = TypedDict("TD", {"key": int, "value": T})
+reveal_type(TD)  # N: Revealed type is "def [T] (*, key: builtins.int, value: T`-1) -> TypedDict('__main__.TD', {'key': builtins.int, 'value': T`-1})"
+
+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")
+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]