Skip to content

Commit 0cf7ed7

Browse files
ilevkivskyiilevkivskyi
authored
Add basic support for PEP 798 (#21976)
Fixes #21489 Things are mostly straightforward. I also add support for old parser. Two non-trivial things: * I wrap constructor arguments instead of using `ARG_STAR` actual kind, because there is a limitation in `argmap.py` that it can't handle structural iterables. We have a bunch of issues open about this, and there were some attempts to fix it, but they stalled. * I force union inference for star-unpacks in comprehensions. Otherwise we will get too many `builtins.object` if we use joins. I think it is fine, since this is a new syntax.
1 parent 6632e40 commit 0cf7ed7

9 files changed

Lines changed: 107 additions & 16 deletions

File tree

‎mypy/checkexpr.py‎

Lines changed: 47 additions & 11 deletions
Original file line numberDiff line numberDiff line change
@@ -5980,17 +5980,39 @@ def check_generator_or_comprehension(
59805980
upper_bound=self.object_type(),
59815981
default=AnyType(TypeOfAny.from_omitted_generics),
59825982
)
5983+
if isinstance(gen.left_expr, StarExpr):
5984+
left_expr = gen.left_expr.expr
5985+
# Note: we wrap the argument type instead of using ARG_STAR kind, because logic
5986+
# in argmap.py doesn't work with structural subtypes of Iterable.
5987+
arg = self.chk.named_generic_type("typing.Iterable", [tv])
5988+
# Motivation for inferring more unions in this case is two-fold:
5989+
# * In regular (non-star) case the constructor signature has bare type variable,
5990+
# thus hitting a special case in constraints solver that would naturally infer
5991+
# more unions. We want to match this behavior here.
5992+
# * This is a new syntax, so we can experiment with something
5993+
# we may want in long-term
5994+
force_infer_unions = True
5995+
else:
5996+
left_expr = gen.left_expr
5997+
arg = tv
5998+
force_infer_unions = False
59835999
tv_list: list[Type] = [tv]
59846000
constructor = CallableType(
5985-
tv_list,
6001+
[arg],
59866002
[nodes.ARG_POS],
59876003
[None],
59886004
self.chk.named_generic_type(type_name, tv_list + additional_args),
59896005
self.chk.named_type("builtins.function"),
59906006
name=id_for_messages,
59916007
variables=[tv],
59926008
)
5993-
return self.check_call(constructor, [gen.left_expr], [nodes.ARG_POS], gen)[0]
6009+
if force_infer_unions:
6010+
old_infer_unions = type_state.infer_unions
6011+
type_state.infer_unions = True
6012+
res = self.check_call(constructor, [left_expr], [nodes.ARG_POS], gen)
6013+
if force_infer_unions:
6014+
type_state.infer_unions = old_infer_unions
6015+
return res[0]
59946016

59956017
def visit_dictionary_comprehension(self, e: DictionaryComprehension) -> Type:
59966018
"""Type check a dictionary comprehension."""
@@ -6015,21 +6037,35 @@ def visit_dictionary_comprehension(self, e: DictionaryComprehension) -> Type:
60156037
upper_bound=self.object_type(),
60166038
default=AnyType(TypeOfAny.from_omitted_generics),
60176039
)
6040+
if e.key is None:
6041+
# Logic and motivation here is similar to check_generator_or_comprehension().
6042+
arg_types = [self.chk.named_generic_type("typing.Mapping", [ktdef, vtdef])]
6043+
arg_kinds = [nodes.ARG_POS]
6044+
arg_names = [None]
6045+
args = [e.value]
6046+
force_infer_unions = True
6047+
else:
6048+
arg_types = [ktdef, vtdef]
6049+
arg_kinds = [nodes.ARG_POS, nodes.ARG_POS]
6050+
arg_names = [None, None]
6051+
args = [e.key, e.value]
6052+
force_infer_unions = False
60186053
constructor = CallableType(
6019-
[ktdef, vtdef],
6020-
[nodes.ARG_POS, nodes.ARG_POS],
6021-
[None, None],
6054+
arg_types,
6055+
arg_kinds,
6056+
arg_names,
60226057
self.chk.named_generic_type("builtins.dict", [ktdef, vtdef]),
60236058
self.chk.named_type("builtins.function"),
60246059
name="<dictionary-comprehension>",
60256060
variables=[ktdef, vtdef],
60266061
)
6027-
if e.key is None:
6028-
self.chk.fail("PEP 798 is not supported yet", e)
6029-
return AnyType(TypeOfAny.from_error)
6030-
return self.check_call(
6031-
constructor, [e.key, e.value], [nodes.ARG_POS, nodes.ARG_POS], e
6032-
)[0]
6062+
if force_infer_unions:
6063+
old_infer_unions = type_state.infer_unions
6064+
type_state.infer_unions = True
6065+
res = self.check_call(constructor, args, arg_kinds, e)
6066+
if force_infer_unions:
6067+
type_state.infer_unions = old_infer_unions
6068+
return res[0]
60336069

60346070
def check_for_comp(self, e: GeneratorExpr | DictionaryComprehension) -> None:
60356071
"""Check the for_comp part of comprehensions. That is the part from 'for':

‎mypy/fastparse.py‎

Lines changed: 9 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -1569,9 +1569,15 @@ def visit_DictComp(self, n: ast3.DictComp) -> DictionaryComprehension:
15691569
iters = [self.visit(c.iter) for c in n.generators]
15701570
ifs_list = [self.translate_expr_list(c.ifs) for c in n.generators]
15711571
is_async = [bool(c.is_async) for c in n.generators]
1572-
e = DictionaryComprehension(
1573-
self.visit(n.key), self.visit(n.value), targets, iters, ifs_list, is_async
1574-
)
1572+
key = self.visit(n.key)
1573+
if n.value is not None:
1574+
value = self.visit(n.value)
1575+
else:
1576+
# Our convention for ** unpack in dictionary comprehension matches Ruff
1577+
# parser, while Python parser has an opposite one.
1578+
value = key # type: ignore[unreachable]
1579+
key = None
1580+
e = DictionaryComprehension(key, value, targets, iters, ifs_list, is_async)
15751581
return self.set_line(e, n)
15761582

15771583
# GeneratorExp(expr elt, comprehension* generators)

‎mypy/nativeparse.py‎

Lines changed: 4 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1620,7 +1620,6 @@ def read_expression(state: State, data: ReadBuffer) -> Expression:
16201620
expr = DictionaryComprehension(key, value, indices, sequences, condlists, is_async)
16211621
read_loc(data, expr)
16221622
if key is None:
1623-
# TODO: add similar check to other kinds of comprehensions.
16241623
state.check_min_version("Unpacking in comprehensions", (3, 15), expr.line, expr.column)
16251624
expect_end_tag(data)
16261625
return expr
@@ -1804,6 +1803,10 @@ def read_expression_list(state: State, data: ReadBuffer) -> list[Expression]:
18041803
def read_generator_expr(state: State, data: ReadBuffer) -> GeneratorExpr:
18051804
"""Helper function to read comprehension data (shared by Generator, ListComp, SetComp)"""
18061805
left_expr = read_expression(state, data)
1806+
if isinstance(left_expr, StarExpr):
1807+
state.check_min_version(
1808+
"Unpacking in comprehensions", (3, 15), left_expr.line, left_expr.column
1809+
)
18071810
n_generators = read_int(data)
18081811
indices = [read_expression(state, data) for _ in range(n_generators)]
18091812
sequences = [read_expression(state, data) for _ in range(n_generators)]

‎mypy/semanal.py‎

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -6493,6 +6493,8 @@ def visit_dictionary_comprehension(self, expr: DictionaryComprehension) -> None:
64936493
self.analyze_comp_for_2(expr)
64946494

64956495
def visit_generator_expr(self, expr: GeneratorExpr) -> None:
6496+
if isinstance(expr.left_expr, StarExpr):
6497+
expr.left_expr.valid = True
64966498
with self.enter(expr):
64976499
self.analyze_comp_for(expr)
64986500
expr.left_expr.accept(self)

‎mypy/test/testcheck.py‎

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -52,6 +52,8 @@
5252
typecheck_files.remove("check-python313.test")
5353
if sys.version_info < (3, 14):
5454
typecheck_files.remove("check-python314.test")
55+
if sys.version_info < (3, 15):
56+
typecheck_files.remove("check-python315.test")
5557

5658

5759
class TypeCheckSuite(DataSuite):

‎mypyc/irbuild/expression.py‎

Lines changed: 10 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1390,6 +1390,9 @@ def _visit_display(
13901390

13911391

13921392
def transform_list_comprehension(builder: IRBuilder, o: ListComprehension) -> Value:
1393+
if isinstance(o.generator.left_expr, StarExpr):
1394+
builder.error("Unpacking in comprehensions is not supported yet", o.line)
1395+
return builder.none()
13931396
gen = o.generator
13941397
if gen in builder.comprehension_to_fitem:
13951398
return _translate_comprehension_with_scope(
@@ -1399,6 +1402,9 @@ def transform_list_comprehension(builder: IRBuilder, o: ListComprehension) -> Va
13991402

14001403

14011404
def transform_set_comprehension(builder: IRBuilder, o: SetComprehension) -> Value:
1405+
if isinstance(o.generator.left_expr, StarExpr):
1406+
builder.error("Unpacking in comprehensions is not supported yet", o.line)
1407+
return builder.none()
14021408
gen = o.generator
14031409
if gen in builder.comprehension_to_fitem:
14041410
return _translate_comprehension_with_scope(
@@ -1424,7 +1430,7 @@ def gen_inner_stmts() -> None:
14241430
if o.key is not None:
14251431
k = builder.accept(o.key)
14261432
else:
1427-
builder.error("PEP 798 is not supported yet", o.line)
1433+
builder.error("Unpacking in comprehensions is not supported yet", o.line)
14281434
k = builder.none()
14291435
v = builder.accept(o.value)
14301436
builder.call_c(exact_dict_set_item_op, [builder.read(d, o.line), k, v], o.line)
@@ -1473,6 +1479,9 @@ def get_arg(arg: Expression | None) -> Value:
14731479

14741480

14751481
def transform_generator_expr(builder: IRBuilder, o: GeneratorExpr) -> Value:
1482+
if isinstance(o.left_expr, StarExpr):
1483+
builder.error("Unpacking in comprehensions is not supported yet", o.line)
1484+
return builder.none()
14761485
builder.warning("Treating generator comprehension as list", o.line)
14771486
if o in builder.comprehension_to_fitem:
14781487
return builder.primitive_op(

‎mypyc/test-data/irbuild-python315.test‎

Lines changed: 6 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,6 @@
1+
[case testPEP798FailGracefully_python3_15]
2+
x = [[1]]
3+
y = [*l for l in x] # E: Unpacking in comprehensions is not supported yet
4+
5+
dx = [{1: 1}]
6+
dy = {**d for d in dx} # E: Unpacking in comprehensions is not supported yet

‎mypyc/test/test_irbuild.py‎

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -73,6 +73,9 @@
7373
if sys.version_info >= (3, 14):
7474
files.append("irbuild-python314.test")
7575

76+
if sys.version_info >= (3, 15):
77+
files.append("irbuild-python315.test")
78+
7679

7780
class TestGenOps(MypycDataSuite):
7881
files = files

‎test-data/unit/check-python315.test‎

Lines changed: 24 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,24 @@
1+
[case testStarUnpackInComprehensionBasicsNative]
2+
# flags: --native-parser
3+
4+
def test(x: list[list[int]], y: list[dict[str, int]]) -> None:
5+
reveal_type([*i for i in x]) # N: Revealed type is "builtins.list[builtins.int]"
6+
reveal_type({*i for i in x}) # N: Revealed type is "builtins.set[builtins.int]"
7+
reveal_type({**d for d in y}) # N: Revealed type is "builtins.dict[builtins.str, builtins.int]"
8+
[builtins fixtures/isinstancelist.pyi]
9+
10+
[case testStarUnpackInComprehensionBasicsOld]
11+
# flags: --no-native-parser
12+
13+
def test(x: list[list[int]], y: list[dict[str, int]]) -> None:
14+
reveal_type([*i for i in x]) # N: Revealed type is "builtins.list[builtins.int]"
15+
reveal_type({*i for i in x}) # N: Revealed type is "builtins.set[builtins.int]"
16+
reveal_type({**d for d in y}) # N: Revealed type is "builtins.dict[builtins.str, builtins.int]"
17+
[builtins fixtures/isinstancelist.pyi]
18+
19+
[case testStarUnpackInComprehensionUnion]
20+
def test(x: list[list[int]] | list[list[str]], y: list[dict[str, int]] | list[dict[str, str]]) -> None:
21+
reveal_type([*i for i in x]) # N: Revealed type is "builtins.list[builtins.int | builtins.str]"
22+
reveal_type({*i for i in x}) # N: Revealed type is "builtins.set[builtins.int | builtins.str]"
23+
reveal_type({**d for d in y}) # N: Revealed type is "builtins.dict[builtins.str | builtins.str, builtins.int | builtins.str]"
24+
[builtins fixtures/isinstancelist.pyi]

0 commit comments

Comments
 (0)