From 15ab6ffbbafd732701158cb9cfbd7ff189377563 Mon Sep 17 00:00:00 2001 From: Matthew Mckee Date: Wed, 2 Apr 2025 13:24:59 +0100 Subject: [PATCH 1/8] [red-knot] Type narrowing for assertions --- .../resources/mdtest/narrow/assert.md | 9 +++++++++ .../src/semantic_index/builder.rs | 11 +++++++++++ crates/red_knot_python_semantic/src/types/infer.rs | 2 +- 3 files changed, 21 insertions(+), 1 deletion(-) create mode 100644 crates/red_knot_python_semantic/resources/mdtest/narrow/assert.md diff --git a/crates/red_knot_python_semantic/resources/mdtest/narrow/assert.md b/crates/red_knot_python_semantic/resources/mdtest/narrow/assert.md new file mode 100644 index 00000000000000..8c152968d5077d --- /dev/null +++ b/crates/red_knot_python_semantic/resources/mdtest/narrow/assert.md @@ -0,0 +1,9 @@ +# Narrowing with assert statements + +## `assert` on singletons + +```py +def _(x: str | None): + assert x is not None + reveal_type(x) # revealed: str +``` diff --git a/crates/red_knot_python_semantic/src/semantic_index/builder.rs b/crates/red_knot_python_semantic/src/semantic_index/builder.rs index af4c4b83c60e55..47430378cf94c4 100644 --- a/crates/red_knot_python_semantic/src/semantic_index/builder.rs +++ b/crates/red_knot_python_semantic/src/semantic_index/builder.rs @@ -1292,6 +1292,17 @@ where ); } } + + ast::Stmt::Assert(node) => { + self.visit_expr(&node.test); + let no_branch_taken = self.flow_snapshot(); + let last_predicate = self.record_expression_narrowing_constraint(&node.test); + + self.record_visibility_constraint(last_predicate); + + self.simplify_visibility_constraints(no_branch_taken); + } + ast::Stmt::Assign(node) => { debug_assert_eq!(&self.current_assignments, &[]); diff --git a/crates/red_knot_python_semantic/src/types/infer.rs b/crates/red_knot_python_semantic/src/types/infer.rs index 8f8d9b2f46e3f4..f81a156b3910f1 100644 --- a/crates/red_knot_python_semantic/src/types/infer.rs +++ b/crates/red_knot_python_semantic/src/types/infer.rs @@ -3188,7 +3188,7 @@ impl<'db> TypeInferenceBuilder<'db> { msg, } = assert; - let test_ty = self.infer_expression(test); + let test_ty = self.infer_standalone_expression(test); if let Err(err) = test_ty.try_bool(self.db()) { err.report_diagnostic(&self.context, &**test); From 800f1fc4d0dd42c6caeedcca7a5109fcb013bdd1 Mon Sep 17 00:00:00 2001 From: Matthew Mckee Date: Wed, 2 Apr 2025 13:48:33 +0100 Subject: [PATCH 2/8] Fix tests --- .../resources/mdtest/narrow/assert.md | 26 +++++++++++++++++-- .../src/semantic_index/builder.rs | 5 +++- 2 files changed, 28 insertions(+), 3 deletions(-) diff --git a/crates/red_knot_python_semantic/resources/mdtest/narrow/assert.md b/crates/red_knot_python_semantic/resources/mdtest/narrow/assert.md index 8c152968d5077d..e7c1b1684767a0 100644 --- a/crates/red_knot_python_semantic/resources/mdtest/narrow/assert.md +++ b/crates/red_knot_python_semantic/resources/mdtest/narrow/assert.md @@ -1,9 +1,31 @@ # Narrowing with assert statements -## `assert` on singletons +## `assert` with None ```py -def _(x: str | None): +def _(x: str | None, y: str | None): assert x is not None reveal_type(x) # revealed: str + assert y is None + reveal_type(y) # revealed: None +``` + +## `assert` on bool + +```py +def _(x: bool, y: bool): + assert x + reveal_type(x) # revealed: Literal[True] + assert not y + reveal_type(y) # revealed: Literal[False] +``` + +## `assert` on `Literal` + +```py +from typing import Literal + +def _(x: Literal[1, 2, 3]): + assert x == 4 + reveal_type(x) # revealed: Never ``` diff --git a/crates/red_knot_python_semantic/src/semantic_index/builder.rs b/crates/red_knot_python_semantic/src/semantic_index/builder.rs index 47430378cf94c4..d3096acb8913b2 100644 --- a/crates/red_knot_python_semantic/src/semantic_index/builder.rs +++ b/crates/red_knot_python_semantic/src/semantic_index/builder.rs @@ -1297,9 +1297,12 @@ where self.visit_expr(&node.test); let no_branch_taken = self.flow_snapshot(); let last_predicate = self.record_expression_narrowing_constraint(&node.test); - self.record_visibility_constraint(last_predicate); + if let Some(msg) = &node.msg { + self.visit_expr(msg); + } + self.simplify_visibility_constraints(no_branch_taken); } From 453602fb536a82c958169ce7929518026eba9af3 Mon Sep 17 00:00:00 2001 From: Matthew Mckee Date: Wed, 2 Apr 2025 13:49:09 +0100 Subject: [PATCH 3/8] Fix mdtest --- .../resources/mdtest/narrow/assert.md | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/crates/red_knot_python_semantic/resources/mdtest/narrow/assert.md b/crates/red_knot_python_semantic/resources/mdtest/narrow/assert.md index e7c1b1684767a0..c28e31ef5c6ecb 100644 --- a/crates/red_knot_python_semantic/resources/mdtest/narrow/assert.md +++ b/crates/red_knot_python_semantic/resources/mdtest/narrow/assert.md @@ -26,6 +26,6 @@ def _(x: bool, y: bool): from typing import Literal def _(x: Literal[1, 2, 3]): - assert x == 4 - reveal_type(x) # revealed: Never + assert x == 2 + reveal_type(x) # revealed: Literal[2] ``` From 59ec01c2f133e9b4f55a400f349d101c1c29a2cd Mon Sep 17 00:00:00 2001 From: Matthew Mckee Date: Wed, 2 Apr 2025 13:49:48 +0100 Subject: [PATCH 4/8] Fix mdtest --- .../red_knot_python_semantic/resources/mdtest/narrow/assert.md | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/crates/red_knot_python_semantic/resources/mdtest/narrow/assert.md b/crates/red_knot_python_semantic/resources/mdtest/narrow/assert.md index c28e31ef5c6ecb..7c1ff1391725c5 100644 --- a/crates/red_knot_python_semantic/resources/mdtest/narrow/assert.md +++ b/crates/red_knot_python_semantic/resources/mdtest/narrow/assert.md @@ -26,6 +26,6 @@ def _(x: bool, y: bool): from typing import Literal def _(x: Literal[1, 2, 3]): - assert x == 2 + assert x is 2 reveal_type(x) # revealed: Literal[2] ``` From 13dbd7d731d66960f9ce14e3c515c353fdec2fdb Mon Sep 17 00:00:00 2001 From: Matthew Mckee Date: Fri, 4 Apr 2025 00:44:34 +0100 Subject: [PATCH 5/8] Add more tests --- .../resources/mdtest/narrow/assert.md | 26 ++++++++++++++++--- 1 file changed, 23 insertions(+), 3 deletions(-) diff --git a/crates/red_knot_python_semantic/resources/mdtest/narrow/assert.md b/crates/red_knot_python_semantic/resources/mdtest/narrow/assert.md index 7c1ff1391725c5..f6add5165a2fe5 100644 --- a/crates/red_knot_python_semantic/resources/mdtest/narrow/assert.md +++ b/crates/red_knot_python_semantic/resources/mdtest/narrow/assert.md @@ -1,6 +1,6 @@ # Narrowing with assert statements -## `assert` with None +## `assert` a value `is None` or `is not None` ```py def _(x: str | None, y: str | None): @@ -10,7 +10,7 @@ def _(x: str | None, y: str | None): reveal_type(y) # revealed: None ``` -## `assert` on bool +## `assert` a value is truthy or falsy ```py def _(x: bool, y: bool): @@ -20,7 +20,7 @@ def _(x: bool, y: bool): reveal_type(y) # revealed: Literal[False] ``` -## `assert` on `Literal` +## `assert` a value `is` a literal ```py from typing import Literal @@ -29,3 +29,23 @@ def _(x: Literal[1, 2, 3]): assert x is 2 reveal_type(x) # revealed: Literal[2] ``` + +## `assert` with `isinstance` + +```py +def _(x: int | str): + assert isinstance(x, int) + reveal_type(x) # revealed: int +``` + +## `assert` a value `in` a tuple + +```py +from typing import Literal + +def _(x: Literal[1, 2, 3], y: Literal[1, 2, 3]): + assert x in (1, 2) + reveal_type(x) # revealed: Literal[1, 2] + assert y not in (1, 2) + reveal_type(y) # revealed: Literal[3] +``` From d732f2ae54e427a775f127bd7d3d8345dea75301 Mon Sep 17 00:00:00 2001 From: Matthew Mckee Date: Sat, 5 Apr 2025 19:23:40 +0100 Subject: [PATCH 6/8] Update test --- .../resources/mdtest/narrow/assert.md | 6 ++++-- 1 file changed, 4 insertions(+), 2 deletions(-) diff --git a/crates/red_knot_python_semantic/resources/mdtest/narrow/assert.md b/crates/red_knot_python_semantic/resources/mdtest/narrow/assert.md index f6add5165a2fe5..1fad9290a51155 100644 --- a/crates/red_knot_python_semantic/resources/mdtest/narrow/assert.md +++ b/crates/red_knot_python_semantic/resources/mdtest/narrow/assert.md @@ -20,14 +20,16 @@ def _(x: bool, y: bool): reveal_type(y) # revealed: Literal[False] ``` -## `assert` a value `is` a literal +## `assert` with `is` and `==` for literals ```py from typing import Literal -def _(x: Literal[1, 2, 3]): +def _(x: Literal[1, 2, 3], y: Literal[1, 2, 3]): assert x is 2 reveal_type(x) # revealed: Literal[2] + assert y == 2 + reveal_type(y) # revealed: Literal[1, 2, 3] ``` ## `assert` with `isinstance` From e409b59e8a6642d5300b37b651a0546a3a780be4 Mon Sep 17 00:00:00 2001 From: Matthew Mckee Date: Sat, 5 Apr 2025 19:24:01 +0100 Subject: [PATCH 7/8] Update crates/red_knot_python_semantic/src/semantic_index/builder.rs Co-authored-by: Carl Meyer --- .../red_knot_python_semantic/src/semantic_index/builder.rs | 7 ++----- 1 file changed, 2 insertions(+), 5 deletions(-) diff --git a/crates/red_knot_python_semantic/src/semantic_index/builder.rs b/crates/red_knot_python_semantic/src/semantic_index/builder.rs index d3096acb8913b2..75b1428b9e2ad7 100644 --- a/crates/red_knot_python_semantic/src/semantic_index/builder.rs +++ b/crates/red_knot_python_semantic/src/semantic_index/builder.rs @@ -1295,15 +1295,12 @@ where ast::Stmt::Assert(node) => { self.visit_expr(&node.test); - let no_branch_taken = self.flow_snapshot(); - let last_predicate = self.record_expression_narrowing_constraint(&node.test); - self.record_visibility_constraint(last_predicate); + let predicate = self.record_expression_narrowing_constraint(&node.test); + self.record_visibility_constraint(predicate); if let Some(msg) = &node.msg { self.visit_expr(msg); } - - self.simplify_visibility_constraints(no_branch_taken); } ast::Stmt::Assign(node) => { From 488c6f7aad5cfcfaebd8ec9f22686b122704a27b Mon Sep 17 00:00:00 2001 From: Carl Meyer Date: Thu, 10 Apr 2025 00:53:34 -0400 Subject: [PATCH 8/8] clippy --- crates/red_knot_python_semantic/src/semantic_index/builder.rs | 1 - 1 file changed, 1 deletion(-) diff --git a/crates/red_knot_python_semantic/src/semantic_index/builder.rs b/crates/red_knot_python_semantic/src/semantic_index/builder.rs index 75b1428b9e2ad7..9db4ea1b9aaf22 100644 --- a/crates/red_knot_python_semantic/src/semantic_index/builder.rs +++ b/crates/red_knot_python_semantic/src/semantic_index/builder.rs @@ -534,7 +534,6 @@ impl<'db> SemanticIndexBuilder<'db> { } /// Records a visibility constraint by applying it to all live bindings and declarations. - #[must_use = "A visibility constraint must always be negated after it is added"] fn record_visibility_constraint( &mut self, predicate: Predicate<'db>,