Add support for mapping to supertype

ilevkivskyi · Aug 15, 2022 · Aug 12, 2022 · Aug 12, 2022 · Aug 12, 2022 · Aug 12, 2022
commit 83f78084197d038a466031b73a50bc228c43def9
diff --git a/mypy/maptype.py b/mypy/maptype.py
@@ -1,8 +1,19 @@
 from typing import Dict, List
 
+import mypy.typeops
 from mypy.expandtype import expand_type
 from mypy.nodes import TypeInfo
-from mypy.types import AnyType, Instance, ProperType, Type, TypeOfAny, TypeVarId
+from mypy.types import (
+    AnyType,
+    Instance,
+    ProperType,
+    TupleType,
+    Type,
+    TypeOfAny,
+    TypeVarId,
+    get_proper_type,
+    has_type_vars,
+)
 
 
 def map_instance_to_supertype(instance: Instance, superclass: TypeInfo) -> Instance:
@@ -16,6 +27,20 @@ def map_instance_to_supertype(instance: Instance, superclass: TypeInfo) -> Insta
         # Fast path: `instance` already belongs to `superclass`.
         return instance
 
+    if superclass.fullname == "builtins.tuple" and instance.type.tuple_type:
+        if has_type_vars(instance.type.tuple_type):
+            # We special case mapping generic tuple types to tuple base, because for
+            # such tuples fallback can't be calculated before applying type arguments.
+            alias = instance.type.special_alias
+            assert alias is not None
+            if not alias._is_recursive:
+                # Unfortunately we can't support this for generic recursive tuples.
+                # If we skip this special casing we will fall back to tuple[Any, ...].
+                env = instance_to_type_environment(instance)
+                tuple_type = get_proper_type(expand_type(instance.type.tuple_type, env))
+                if isinstance(tuple_type, TupleType):
+                    return mypy.typeops.tuple_fallback(tuple_type)
+
     if not superclass.type_vars:
         # Fast path: `superclass` has no type variables to map to.
         return Instance(superclass, [])

diff --git a/mypy/semanal_namedtuple.py b/mypy/semanal_namedtuple.py
@@ -58,6 +58,7 @@
     TypeType,
     TypeVarType,
     UnboundType,
+    has_type_vars,
 )
 from mypy.util import get_unique_redefinition_name
 
@@ -487,7 +488,7 @@ def build_namedtuple_typeinfo(
         # We can't calculate the complete fallback type until after semantic
         # analysis, since otherwise base classes might be incomplete. Postpone a
         # callback function that patches the fallback.
-        if not has_placeholder(tuple_base):
+        if not has_placeholder(tuple_base) and not has_type_vars(tuple_base):
             self.api.schedule_patch(
                 PRIORITY_FALLBACKS, lambda: calculate_tuple_fallback(tuple_base)
             )

diff --git a/test-data/unit/check-namedtuple.test b/test-data/unit/check-namedtuple.test
@@ -1186,23 +1186,79 @@ NT[str](key=0, value=0)  # E: Argument "value" to "NT" has incompatible type "in
 [case testGenericNamedTupleMethods]
 from typing import Generic, NamedTuple, TypeVar
 
+T = TypeVar("T")
+class NT(NamedTuple, Generic[T]):
+    key: int
+    value: T
+x: int
+
+nti: NT[int]
+reveal_type(nti * x)  # N: Revealed type is "builtins.tuple[builtins.int, ...]"
+
+nts: NT[str]
+reveal_type(nts * x)  # N: Revealed type is "builtins.tuple[builtins.object, ...]"
 [builtins fixtures/tuple.pyi]
 [typing fixtures/typing-namedtuple.pyi]
 
 [case testGenericNamedTupleCustomMethods]
 from typing import Generic, NamedTuple, TypeVar
 
+T = TypeVar("T")
+class NT(NamedTuple, Generic[T]):
+    key: int
+    value: T
+    def foo(self) -> T: ...
+    @classmethod
+    def from_value(cls, value: T) -> NT[T]: ...
+
+nts: NT[str]
+reveal_type(nts.foo())  # N: Revealed type is "builtins.str"
+
+nti = NT.from_value(1)
+reveal_type(nti)  # N: Revealed type is "Tuple[builtins.int, builtins.int, fallback=__main__.NT[builtins.int]]"
+NT[str].from_value(1)  # E: Argument 1 to "from_value" of "NT" has incompatible type "int"; expected "str"
 [builtins fixtures/tuple.pyi]
 [typing fixtures/typing-namedtuple.pyi]
 
 [case testGenericNamedTupleSubtyping]
-from typing import Generic, NamedTuple, TypeVar
+from typing import Generic, NamedTuple, TypeVar, Tuple
 
+T = TypeVar("T")
+class NT(NamedTuple, Generic[T]):
+    key: int
+    value: T
+
+nts: NT[str]
+nti: NT[int]
+
+def foo(x: Tuple[int, ...]) -> None: ...
+foo(nti)
+foo(nts)  # E: Argument 1 to "foo" has incompatible type "NT[str]"; expected "Tuple[int, ...]"
 [builtins fixtures/tuple.pyi]
 [typing fixtures/typing-namedtuple.pyi]
 
 [case testGenericNamedTupleJoin]
-from typing import Generic, NamedTuple, TypeVar
+from typing import Generic, NamedTuple, TypeVar, Tuple
+
+T = TypeVar("T")
+class NT(NamedTuple, Generic[T]):
+    key: int
+    value: T
+
+nts: NT[str]
+nti: NT[int]
+x: Tuple[int, ...]
+
+def foo(x: T, y: T) -> T: ...
+reveal_type(foo(nti, nti))  # N: Revealed type is "Tuple[builtins.int, builtins.int, fallback=__main__.NT[builtins.int]]"
+
+# Here fallbacks are object because named tuples are invariant.
+reveal_type(foo(nti, nts))  # N: Revealed type is "Tuple[builtins.int, builtins.object, fallback=builtins.object]"
+reveal_type(foo(nts, nti))  # N: Revealed type is "Tuple[builtins.int, builtins.object, fallback=builtins.object]"
 
+reveal_type(foo(nti, x))  # N: Revealed type is "builtins.tuple[builtins.int, ...]"
+reveal_type(foo(nts, x))  # N: Revealed type is "builtins.tuple[builtins.object, ...]"
+reveal_type(foo(x, nti))  # N: Revealed type is "builtins.tuple[builtins.int, ...]"
+reveal_type(foo(x, nts))  # N: Revealed type is "builtins.tuple[builtins.object, ...]"
 [builtins fixtures/tuple.pyi]
 [typing fixtures/typing-namedtuple.pyi]
diff --git a/test-data/unit/fixtures/tuple.pyi b/test-data/unit/fixtures/tuple.pyi
@@ -25,6 +25,7 @@ class tuple(Sequence[Tco], Generic[Tco]):
     def count(self, obj: object) -> int: pass
 class function: pass
 class ellipsis: pass
+class classmethod: pass
 
 # We need int and slice for indexing tuples.
 class int: