diff --git a/crates/ty_python_semantic/resources/mdtest/narrow/conditionals/nested.md b/crates/ty_python_semantic/resources/mdtest/narrow/conditionals/nested.md index 033ca89d3e..654e158d81 100644 --- a/crates/ty_python_semantic/resources/mdtest/narrow/conditionals/nested.md +++ b/crates/ty_python_semantic/resources/mdtest/narrow/conditionals/nested.md @@ -43,6 +43,23 @@ def _(flag1: bool, flag2: bool): reveal_type(x) # revealed: Never ``` +## Comprehensions + +```py +def _(xs: list[int | None], ys: list[str | bytes], list_of_optional_lists: list[list[int | None] | None]): + [reveal_type(x) for x in xs if x is not None] # revealed: int + [reveal_type(y) for y in ys if isinstance(y, str)] # revealed: str + + [_ for x in xs if x is not None if reveal_type(x) // 3 != 0] # revealed: int + + [reveal_type(x) for x in xs if x is not None if x != 0 if x != 1] # revealed: int & ~Literal[0] & ~Literal[1] + + [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] + [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] + + [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 +``` + ## Cross-scope narrowing Narrowing constraints are also valid in eager nested scopes (however, because class variables are @@ -194,9 +211,7 @@ def f(x: str | None): if l[0] is not None: reveal_type(l[0]) # revealed: str - # TODO: should be str - # This could be fixed if we supported narrowing with if clauses in comprehensions. - [reveal_type(x) for _ in range(1) if x is not None] # revealed: str | None + [reveal_type(x) for _ in range(1) if x is not None] # revealed: str ``` ### Narrowing constraints introduced in the outer scope @@ -276,8 +291,7 @@ def f(x: str | Literal[1] | None): if x != 1: reveal_type(x) # revealed: str - # TODO: should be str - [reveal_type(x) for _ in range(1) if x != 1] # revealed: str | Literal[1] + [reveal_type(x) for _ in range(1) if x != 1] # revealed: str if g is not None: def _(): diff --git a/crates/ty_python_semantic/src/semantic_index/builder.rs b/crates/ty_python_semantic/src/semantic_index/builder.rs index 1fc3e3dd1f..dcaff07143 100644 --- a/crates/ty_python_semantic/src/semantic_index/builder.rs +++ b/crates/ty_python_semantic/src/semantic_index/builder.rs @@ -854,8 +854,9 @@ impl<'db, 'ast> SemanticIndexBuilder<'db, 'ast> { value, ); - for expr in &generator.ifs { - self.visit_expr(expr); + for if_expr in &generator.ifs { + self.visit_expr(if_expr); + self.record_expression_narrowing_constraint(if_expr); } for generator in generators_iter { @@ -871,8 +872,9 @@ impl<'db, 'ast> SemanticIndexBuilder<'db, 'ast> { value, ); - for expr in &generator.ifs { - self.visit_expr(expr); + for if_expr in &generator.ifs { + self.visit_expr(if_expr); + self.record_expression_narrowing_constraint(if_expr); } } diff --git a/crates/ty_python_semantic/src/types/infer.rs b/crates/ty_python_semantic/src/types/infer.rs index 0240c01f8a..d3aabea5c6 100644 --- a/crates/ty_python_semantic/src/types/infer.rs +++ b/crates/ty_python_semantic/src/types/infer.rs @@ -5089,7 +5089,7 @@ impl<'db, 'ast> TypeInferenceBuilder<'db, 'ast> { .iterate(builder.db()) }); for expr in ifs { - self.infer_expression(expr); + self.infer_standalone_expression(expr); } }