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 0000000000..1fad9290a5 --- /dev/null +++ b/crates/red_knot_python_semantic/resources/mdtest/narrow/assert.md @@ -0,0 +1,53 @@ +# Narrowing with assert statements + +## `assert` a value `is None` or `is not None` + +```py +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` a value is truthy or falsy + +```py +def _(x: bool, y: bool): + assert x + reveal_type(x) # revealed: Literal[True] + assert not y + reveal_type(y) # revealed: Literal[False] +``` + +## `assert` with `is` and `==` for literals + +```py +from typing import Literal + +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` + +```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] +``` 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 af4c4b83c6..9db4ea1b9a 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>, @@ -1292,6 +1291,17 @@ where ); } } + + ast::Stmt::Assert(node) => { + self.visit_expr(&node.test); + let predicate = self.record_expression_narrowing_constraint(&node.test); + self.record_visibility_constraint(predicate); + + if let Some(msg) = &node.msg { + self.visit_expr(msg); + } + } + 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 8f8d9b2f46..f81a156b39 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);