Enable definition of generic tuple types

python · ilevkivskyi · Aug 15, 2022 · Aug 12, 2022 · Aug 12, 2022 · Aug 12, 2022
commit 550423eddac904d99139a762077cd95bdda11daa
diff --git a/mypy/expandtype.py b/mypy/expandtype.py
@@ -296,7 +296,11 @@ def expand_types_with_unpack(
     def visit_tuple_type(self, t: TupleType) -> Type:
         items = self.expand_types_with_unpack(t.items)
         if isinstance(items, list):
-            return t.copy_modified(items=items)
+            fallback = t.partial_fallback.accept(self)
+            fallback = get_proper_type(fallback)
+            if not isinstance(fallback, Instance):
+                fallback = t.partial_fallback
+            return t.copy_modified(items=items, fallback=fallback)
         else:
             return items
 

diff --git a/mypy/nodes.py b/mypy/nodes.py
@@ -3294,7 +3294,7 @@ def from_tuple_type(cls, info: TypeInfo) -> "TypeAlias":
         """Generate an alias to the tuple type described by a given TypeInfo."""
         assert info.tuple_type
         return TypeAlias(
-            info.tuple_type.copy_modified(fallback=mypy.types.Instance(info, [])),
+            info.tuple_type.copy_modified(fallback=mypy.types.Instance(info, info.defn.type_vars)),
             info.fullname,
             info.line,
             info.column,

diff --git a/mypy/semanal.py b/mypy/semanal.py
@@ -1382,15 +1382,14 @@ def analyze_class(self, defn: ClassDef) -> None:
             return
 
         if self.analyze_namedtuple_classdef(defn):
+            if defn.info:
+                self.setup_type_vars(defn, tvar_defs)
+                self.setup_alias_type_vars(defn)
             return
 
         # Create TypeInfo for class now that base classes and the MRO can be calculated.
         self.prepare_class_def(defn)
-
-        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()
+        self.setup_type_vars(defn, tvar_defs)
         if base_error:
             defn.info.fallback_to_any = True
 
@@ -1403,6 +1402,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)
+        target = defn.info.special_alias.target
+        assert isinstance(target, ProperType) and isinstance(target, TupleType)
+        target.partial_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
 
@@ -1454,7 +1466,7 @@ def analyze_namedtuple_classdef(self, defn: ClassDef) -> bool:
             if info is None:
                 self.mark_incomplete(defn.name, defn)
             else:
-                self.prepare_class_def(defn, info)
+                self.prepare_class_def(defn, info, custom_names=True)
                 with self.scope.class_scope(defn.info):
                     with self.named_tuple_analyzer.save_namedtuple_body(info):
                         self.analyze_class_body_common(defn)
@@ -1679,7 +1691,9 @@ def get_all_bases_tvars(
                 tvars.extend(base_tvars)
         return remove_dups(tvars)
 
-    def prepare_class_def(self, defn: ClassDef, info: Optional[TypeInfo] = None) -> None:
+    def prepare_class_def(
+        self, defn: ClassDef, info: Optional[TypeInfo] = None, custom_names: bool = False
+    ) -> None:
         """Prepare for the analysis of a class definition.
 
         Create an empty TypeInfo and store it in a symbol table, or if the 'info'
@@ -1691,10 +1705,13 @@ def prepare_class_def(self, defn: ClassDef, info: Optional[TypeInfo] = None) ->
             info = info or self.make_empty_type_info(defn)
             defn.info = info
             info.defn = defn
-            if not self.is_func_scope():
-                info._fullname = self.qualified_name(defn.name)
-            else:
-                info._fullname = info.name
+            if not custom_names:
+                # Some special classes (in particular NamedTuples) use custom fullname logic.
+                # Don't override it here (also see comment below, this needs cleanup).
+                if not self.is_func_scope():
+                    info._fullname = self.qualified_name(defn.name)
+                else:
+                    info._fullname = info.name
         local_name = defn.name
         if "@" in local_name:
             local_name = local_name.split("@")[0]
@@ -1855,6 +1872,7 @@ def configure_tuple_base_class(self, defn: ClassDef, base: TupleType) -> Instanc
         if info.special_alias and has_placeholder(info.special_alias.target):
             self.defer(force_progress=True)
         info.update_tuple_type(base)
+        self.setup_alias_type_vars(defn)
 
         if base.partial_fallback.type.fullname == "builtins.tuple" and not has_placeholder(base):
             # Fallback can only be safely calculated after semantic analysis, since base

diff --git a/mypy/semanal_namedtuple.py b/mypy/semanal_namedtuple.py
@@ -116,7 +116,6 @@ def analyze_namedtuple_classdef(
                     info = self.build_namedtuple_typeinfo(
                         defn.name, items, types, default_items, defn.line, existing_info
                     )
-                    defn.info = info
                     defn.analyzed = NamedTupleExpr(info, is_typed=True)
                     defn.analyzed.line = defn.line
                     defn.analyzed.column = defn.column

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

diff --git a/test-data/unit/check-classes.test b/test-data/unit/check-classes.test
@@ -7329,3 +7329,13 @@ a: int = child.foo(1)
 b: str = child.bar("abc")
 c: float = child.baz(3.4)
 d: bool = child.foobar()
+
+[case testGenericTupleTypeCreation]
+from typing import Generic, Tuple, TypeVar
+
+[builtins fixtures/tuple.pyi]
+
+[case testGenericTupleTypeSubclassing]
+from typing import Generic, Tuple, TypeVar
+
+[builtins fixtures/tuple.pyi]
diff --git a/test-data/unit/check-namedtuple.test b/test-data/unit/check-namedtuple.test
@@ -1162,3 +1162,47 @@ NT5 = NamedTuple(b'NT5', [('x', int), ('y', int)])  # E: "NamedTuple()" expects
 
 [builtins fixtures/tuple.pyi]
 [typing fixtures/typing-namedtuple.pyi]
+
+[case testGenericNamedTupleCreation]
+from typing import Generic, NamedTuple, TypeVar
+
+T = TypeVar("T")
+class NT(NamedTuple, Generic[T]):
+    key: int
+    value: T
+
+nts: NT[str]
+reveal_type(nts)  # N: Revealed type is "Tuple[builtins.int, builtins.str, fallback=__main__.NT[builtins.str]]"
+reveal_type(nts.value)  # N: Revealed type is "builtins.str"
+
+nti = NT(key=0, value=0)
+reveal_type(nti)  # N: Revealed type is "Tuple[builtins.int, builtins.int, fallback=__main__.NT[builtins.int]]"
+reveal_type(nti.value)  # N: Revealed type is "builtins.int"
+
+NT[str](key=0, value=0)  # E: Argument "value" to "NT" has incompatible type "int"; expected "str"
+[builtins fixtures/tuple.pyi]
+[typing fixtures/typing-namedtuple.pyi]
+
+[case testGenericNamedTupleMethods]
+from typing import Generic, NamedTuple, TypeVar
+
+[builtins fixtures/tuple.pyi]
+[typing fixtures/typing-namedtuple.pyi]
+
+[case testGenericNamedTupleCustomMethods]
+from typing import Generic, NamedTuple, TypeVar
+
+[builtins fixtures/tuple.pyi]
+[typing fixtures/typing-namedtuple.pyi]
+
+[case testGenericNamedTupleSubtyping]
+from typing import Generic, NamedTuple, TypeVar
+
+[builtins fixtures/tuple.pyi]
+[typing fixtures/typing-namedtuple.pyi]
+
+[case testGenericNamedTupleJoin]
+from typing import Generic, NamedTuple, TypeVar
+
+[builtins fixtures/tuple.pyi]
+[typing fixtures/typing-namedtuple.pyi]
diff --git a/test-data/unit/check-tuples.test b/test-data/unit/check-tuples.test
@@ -883,9 +883,9 @@ from typing import TypeVar, Generic, Tuple
 T = TypeVar('T')
 class Test(Generic[T], Tuple[T]): pass
 x = Test() # type: Test[int]
+reveal_type(x)  # N: Revealed type is "Tuple[builtins.int, fallback=__main__.Test[builtins.int]]"
 [builtins fixtures/tuple.pyi]
 [out]
-main:4: error: Generic tuple types not supported
 
 
 -- Variable-length tuples (Tuple[t, ...] with literal '...')