Skip to content

Commit fbf344b

Browse files
jhancejhance
committed
Implement tuple cases for typevartuple constraints
1 parent 82a97f7 commit fbf344b

File tree

3 files changed

+110
-5
lines changed

3 files changed

+110
-5
lines changed

‎mypy/constraints.py

Lines changed: 13 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -644,11 +644,20 @@ def visit_instance(self, template: Instance) -> list[Constraint]:
644644
isinstance(template_unpack, Instance)
645645
and template_unpack.type.fullname == "builtins.tuple"
646646
):
647-
# TODO: check homogenous tuple case
648-
raise NotImplementedError
647+
for item in mapped_middle:
648+
res.extend(
649+
infer_constraints(
650+
template_unpack.args[0], item, self.direction
651+
)
652+
)
649653
elif isinstance(template_unpack, TupleType):
650-
# TODO: check tuple case
651-
raise NotImplementedError
654+
if len(template_unpack.items) == len(mapped_middle):
655+
for template_arg, item in zip(
656+
template_unpack.items, mapped_middle
657+
):
658+
res.extend(
659+
infer_constraints(template_arg, item, self.direction)
660+
)
652661

653662
mapped_args = mapped_prefix + mapped_suffix
654663
template_args = template_prefix + template_suffix

‎mypy/test/testconstraints.py

Lines changed: 96 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -5,7 +5,7 @@
55
from mypy.constraints import SUBTYPE_OF, SUPERTYPE_OF, Constraint, infer_constraints
66
from mypy.test.helpers import Suite
77
from mypy.test.typefixture import TypeFixture
8-
from mypy.types import Instance, TypeList, UnpackType
8+
from mypy.types import Instance, TupleType, TypeList, UnpackType
99

1010

1111
class ConstraintsSuite(Suite):
@@ -48,3 +48,98 @@ def test_type_var_tuple_with_prefix_and_suffix(self) -> None:
4848
Constraint(type_var=fx.ts, op=SUPERTYPE_OF, target=TypeList([fx.b, fx.c])),
4949
Constraint(type_var=fx.s, op=SUPERTYPE_OF, target=fx.d),
5050
}
51+
52+
def test_unpack_homogenous_tuple(self) -> None:
53+
fx = self.fx
54+
assert set(
55+
infer_constraints(
56+
Instance(fx.gvi, [UnpackType(Instance(fx.std_tuplei, [fx.t]))]),
57+
Instance(fx.gvi, [fx.a, fx.b]),
58+
SUPERTYPE_OF,
59+
)
60+
) == {
61+
Constraint(type_var=fx.t.id, op=SUPERTYPE_OF, target=fx.a),
62+
Constraint(type_var=fx.t.id, op=SUPERTYPE_OF, target=fx.b),
63+
}
64+
65+
def test_unpack_homogenous_tuple_with_prefix_and_suffix(self) -> None:
66+
fx = self.fx
67+
assert set(
68+
infer_constraints(
69+
Instance(fx.gv2i, [fx.t, UnpackType(Instance(fx.std_tuplei, [fx.s])), fx.u]),
70+
Instance(fx.gv2i, [fx.a, fx.b, fx.c, fx.d]),
71+
SUPERTYPE_OF,
72+
)
73+
) == {
74+
Constraint(type_var=fx.t.id, op=SUPERTYPE_OF, target=fx.a),
75+
Constraint(type_var=fx.s.id, op=SUPERTYPE_OF, target=fx.b),
76+
Constraint(type_var=fx.s.id, op=SUPERTYPE_OF, target=fx.c),
77+
Constraint(type_var=fx.u.id, op=SUPERTYPE_OF, target=fx.d),
78+
}
79+
80+
def test_unpack_tuple(self) -> None:
81+
fx = self.fx
82+
assert set(
83+
infer_constraints(
84+
Instance(
85+
fx.gvi,
86+
[
87+
UnpackType(
88+
TupleType([fx.t, fx.s], fallback=Instance(fx.std_tuplei, [fx.o]))
89+
)
90+
],
91+
),
92+
Instance(fx.gvi, [fx.a, fx.b]),
93+
SUPERTYPE_OF,
94+
)
95+
) == {
96+
Constraint(type_var=fx.t.id, op=SUPERTYPE_OF, target=fx.a),
97+
Constraint(type_var=fx.s.id, op=SUPERTYPE_OF, target=fx.b),
98+
}
99+
100+
def test_unpack_with_prefix_and_suffix(self) -> None:
101+
fx = self.fx
102+
assert set(
103+
infer_constraints(
104+
Instance(
105+
fx.gv2i,
106+
[
107+
fx.u,
108+
UnpackType(
109+
TupleType([fx.t, fx.s], fallback=Instance(fx.std_tuplei, [fx.o]))
110+
),
111+
fx.u,
112+
],
113+
),
114+
Instance(fx.gv2i, [fx.a, fx.b, fx.c, fx.d]),
115+
SUPERTYPE_OF,
116+
)
117+
) == {
118+
Constraint(type_var=fx.u.id, op=SUPERTYPE_OF, target=fx.a),
119+
Constraint(type_var=fx.t.id, op=SUPERTYPE_OF, target=fx.b),
120+
Constraint(type_var=fx.s.id, op=SUPERTYPE_OF, target=fx.c),
121+
Constraint(type_var=fx.u.id, op=SUPERTYPE_OF, target=fx.d),
122+
}
123+
124+
def test_unpack_tuple_length_non_match(self) -> None:
125+
fx = self.fx
126+
assert set(
127+
infer_constraints(
128+
Instance(
129+
fx.gv2i,
130+
[
131+
fx.u,
132+
UnpackType(
133+
TupleType([fx.t, fx.s], fallback=Instance(fx.std_tuplei, [fx.o]))
134+
),
135+
fx.u,
136+
],
137+
),
138+
Instance(fx.gv2i, [fx.a, fx.b, fx.d]),
139+
SUPERTYPE_OF,
140+
)
141+
# We still get constraints on the prefix/suffix in this case.
142+
) == {
143+
Constraint(type_var=fx.u.id, op=SUPERTYPE_OF, target=fx.a),
144+
Constraint(type_var=fx.u.id, op=SUPERTYPE_OF, target=fx.d),
145+
}

‎mypy/test/typefixture.py

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -66,6 +66,7 @@ def make_type_var_tuple(name: str, id: int, upper_bound: Type) -> TypeVarTupleTy
6666
self.s1 = make_type_var("S", 1, [], self.o, variance) # S`1 (type variable)
6767
self.sf = make_type_var("S", -2, [], self.o, variance) # S`-2 (type variable)
6868
self.sf1 = make_type_var("S", -1, [], self.o, variance) # S`-1 (type variable)
69+
self.u = make_type_var("U", 3, [], self.o, variance) # U`3 (type variable)
6970

7071
self.ts = make_type_var_tuple("Ts", 1, self.o) # Ts`1 (type var tuple)
7172
self.ss = make_type_var_tuple("Ss", 2, self.o) # Ss`2 (type var tuple)

0 commit comments

Comments
 (0)