Skip to content

Commit 689797a

Browse files
authored
[ty] Type narrowing in comprehensions (#18934)
## Summary Add type narrowing inside comprehensions: ```py def _(xs: list[int | None]): [reveal_type(x) for x in xs if x is not None] # revealed: int ``` closes astral-sh/ty#680 ## Test Plan * New Markdown tests * Made sure the example from astral-sh/ty#680 now checks without errors * Made sure that all removed ecosystem diagnostics were actually false positives
1 parent 66dbea9 commit 689797a

File tree

3 files changed

+26
-10
lines changed

3 files changed

+26
-10
lines changed

crates/ty_python_semantic/resources/mdtest/narrow/conditionals/nested.md

Lines changed: 19 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -43,6 +43,23 @@ def _(flag1: bool, flag2: bool):
4343
reveal_type(x) # revealed: Never
4444
```
4545

46+
## Comprehensions
47+
48+
```py
49+
def _(xs: list[int | None], ys: list[str | bytes], list_of_optional_lists: list[list[int | None] | None]):
50+
[reveal_type(x) for x in xs if x is not None] # revealed: int
51+
[reveal_type(y) for y in ys if isinstance(y, str)] # revealed: str
52+
53+
[_ for x in xs if x is not None if reveal_type(x) // 3 != 0] # revealed: int
54+
55+
[reveal_type(x) for x in xs if x is not None if x != 0 if x != 1] # revealed: int & ~Literal[0] & ~Literal[1]
56+
57+
[reveal_type((x, y)) for x in xs if x is not None for y in ys if isinstance(y, str)] # revealed: tuple[int, str]
58+
[reveal_type((x, y)) for y in ys if isinstance(y, str) for x in xs if x is not None] # revealed: tuple[int, str]
59+
60+
[reveal_type(i) for inner in list_of_optional_lists if inner is not None for i in inner if i is not None] # revealed: int
61+
```
62+
4663
## Cross-scope narrowing
4764

4865
Narrowing constraints are also valid in eager nested scopes (however, because class variables are
@@ -194,9 +211,7 @@ def f(x: str | None):
194211
if l[0] is not None:
195212
reveal_type(l[0]) # revealed: str
196213

197-
# TODO: should be str
198-
# This could be fixed if we supported narrowing with if clauses in comprehensions.
199-
[reveal_type(x) for _ in range(1) if x is not None] # revealed: str | None
214+
[reveal_type(x) for _ in range(1) if x is not None] # revealed: str
200215
```
201216

202217
### Narrowing constraints introduced in the outer scope
@@ -276,8 +291,7 @@ def f(x: str | Literal[1] | None):
276291
if x != 1:
277292
reveal_type(x) # revealed: str
278293

279-
# TODO: should be str
280-
[reveal_type(x) for _ in range(1) if x != 1] # revealed: str | Literal[1]
294+
[reveal_type(x) for _ in range(1) if x != 1] # revealed: str
281295

282296
if g is not None:
283297
def _():

crates/ty_python_semantic/src/semantic_index/builder.rs

Lines changed: 6 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -854,8 +854,9 @@ impl<'db, 'ast> SemanticIndexBuilder<'db, 'ast> {
854854
value,
855855
);
856856

857-
for expr in &generator.ifs {
858-
self.visit_expr(expr);
857+
for if_expr in &generator.ifs {
858+
self.visit_expr(if_expr);
859+
self.record_expression_narrowing_constraint(if_expr);
859860
}
860861

861862
for generator in generators_iter {
@@ -871,8 +872,9 @@ impl<'db, 'ast> SemanticIndexBuilder<'db, 'ast> {
871872
value,
872873
);
873874

874-
for expr in &generator.ifs {
875-
self.visit_expr(expr);
875+
for if_expr in &generator.ifs {
876+
self.visit_expr(if_expr);
877+
self.record_expression_narrowing_constraint(if_expr);
876878
}
877879
}
878880

crates/ty_python_semantic/src/types/infer.rs

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -5089,7 +5089,7 @@ impl<'db, 'ast> TypeInferenceBuilder<'db, 'ast> {
50895089
.iterate(builder.db())
50905090
});
50915091
for expr in ifs {
5092-
self.infer_expression(expr);
5092+
self.infer_standalone_expression(expr);
50935093
}
50945094
}
50955095

0 commit comments

Comments
 (0)