diff --git a/mypy/constraints.py b/mypy/constraints.py index 05bc680230ee..bcbeace1ff2b 100644 --- a/mypy/constraints.py +++ b/mypy/constraints.py @@ -644,11 +644,20 @@ def visit_instance(self, template: Instance) -> list[Constraint]: isinstance(template_unpack, Instance) and template_unpack.type.fullname == "builtins.tuple" ): - # TODO: check homogenous tuple case - raise NotImplementedError + for item in mapped_middle: + res.extend( + infer_constraints( + template_unpack.args[0], item, self.direction + ) + ) elif isinstance(template_unpack, TupleType): - # TODO: check tuple case - raise NotImplementedError + if len(template_unpack.items) == len(mapped_middle): + for template_arg, item in zip( + template_unpack.items, mapped_middle + ): + res.extend( + infer_constraints(template_arg, item, self.direction) + ) mapped_args = mapped_prefix + mapped_suffix template_args = template_prefix + template_suffix diff --git a/mypy/test/testconstraints.py b/mypy/test/testconstraints.py index 4f5d927f956f..6b8f596dd605 100644 --- a/mypy/test/testconstraints.py +++ b/mypy/test/testconstraints.py @@ -5,7 +5,7 @@ from mypy.constraints import SUBTYPE_OF, SUPERTYPE_OF, Constraint, infer_constraints from mypy.test.helpers import Suite from mypy.test.typefixture import TypeFixture -from mypy.types import Instance, TypeList, UnpackType +from mypy.types import Instance, TupleType, TypeList, UnpackType class ConstraintsSuite(Suite): @@ -48,3 +48,98 @@ def test_type_var_tuple_with_prefix_and_suffix(self) -> None: Constraint(type_var=fx.ts, op=SUPERTYPE_OF, target=TypeList([fx.b, fx.c])), Constraint(type_var=fx.s, op=SUPERTYPE_OF, target=fx.d), } + + def test_unpack_homogenous_tuple(self) -> None: + fx = self.fx + assert set( + infer_constraints( + Instance(fx.gvi, [UnpackType(Instance(fx.std_tuplei, [fx.t]))]), + Instance(fx.gvi, [fx.a, fx.b]), + SUPERTYPE_OF, + ) + ) == { + Constraint(type_var=fx.t, op=SUPERTYPE_OF, target=fx.a), + Constraint(type_var=fx.t, op=SUPERTYPE_OF, target=fx.b), + } + + def test_unpack_homogenous_tuple_with_prefix_and_suffix(self) -> None: + fx = self.fx + assert set( + infer_constraints( + Instance(fx.gv2i, [fx.t, UnpackType(Instance(fx.std_tuplei, [fx.s])), fx.u]), + Instance(fx.gv2i, [fx.a, fx.b, fx.c, fx.d]), + SUPERTYPE_OF, + ) + ) == { + Constraint(type_var=fx.t, op=SUPERTYPE_OF, target=fx.a), + Constraint(type_var=fx.s, op=SUPERTYPE_OF, target=fx.b), + Constraint(type_var=fx.s, op=SUPERTYPE_OF, target=fx.c), + Constraint(type_var=fx.u, op=SUPERTYPE_OF, target=fx.d), + } + + def test_unpack_tuple(self) -> None: + fx = self.fx + assert set( + infer_constraints( + Instance( + fx.gvi, + [ + UnpackType( + TupleType([fx.t, fx.s], fallback=Instance(fx.std_tuplei, [fx.o])) + ) + ], + ), + Instance(fx.gvi, [fx.a, fx.b]), + SUPERTYPE_OF, + ) + ) == { + Constraint(type_var=fx.t, op=SUPERTYPE_OF, target=fx.a), + Constraint(type_var=fx.s, op=SUPERTYPE_OF, target=fx.b), + } + + def test_unpack_with_prefix_and_suffix(self) -> None: + fx = self.fx + assert set( + infer_constraints( + Instance( + fx.gv2i, + [ + fx.u, + UnpackType( + TupleType([fx.t, fx.s], fallback=Instance(fx.std_tuplei, [fx.o])) + ), + fx.u, + ], + ), + Instance(fx.gv2i, [fx.a, fx.b, fx.c, fx.d]), + SUPERTYPE_OF, + ) + ) == { + Constraint(type_var=fx.u, op=SUPERTYPE_OF, target=fx.a), + Constraint(type_var=fx.t, op=SUPERTYPE_OF, target=fx.b), + Constraint(type_var=fx.s, op=SUPERTYPE_OF, target=fx.c), + Constraint(type_var=fx.u, op=SUPERTYPE_OF, target=fx.d), + } + + def test_unpack_tuple_length_non_match(self) -> None: + fx = self.fx + assert set( + infer_constraints( + Instance( + fx.gv2i, + [ + fx.u, + UnpackType( + TupleType([fx.t, fx.s], fallback=Instance(fx.std_tuplei, [fx.o])) + ), + fx.u, + ], + ), + Instance(fx.gv2i, [fx.a, fx.b, fx.d]), + SUPERTYPE_OF, + ) + # We still get constraints on the prefix/suffix in this case. + ) == { + Constraint(type_var=fx.u, op=SUPERTYPE_OF, target=fx.a), + Constraint(type_var=fx.u, op=SUPERTYPE_OF, target=fx.d), + } diff --git a/mypy/test/typefixture.py b/mypy/test/typefixture.py index a78ad6e6f51b..380da909893a 100644 --- a/mypy/test/typefixture.py +++ b/mypy/test/typefixture.py @@ -66,6 +66,7 @@ def make_type_var_tuple(name: str, id: int, upper_bound: Type) -> TypeVarTupleTy self.s1 = make_type_var("S", 1, [], self.o, variance) # S`1 (type variable) self.sf = make_type_var("S", -2, [], self.o, variance) # S`-2 (type variable) self.sf1 = make_type_var("S", -1, [], self.o, variance) # S`-1 (type variable) + self.u = make_type_var("U", 3, [], self.o, variance) # U`3 (type variable) self.ts = make_type_var_tuple("Ts", 1, self.o) # Ts`1 (type var tuple) self.ss = make_type_var_tuple("Ss", 2, self.o) # Ss`2 (type var tuple)