Skip to content

Commit 1f7a0af

Browse files
MatthewMckee4carljm
andcommitted
[red-knot] Type narrowing for assertions (#17149)
## Summary Fixes #17147 ## Test Plan Add new narrow/assert.md test file --------- Co-authored-by: Carl Meyer <[email protected]>
1 parent 1a3b737 commit 1f7a0af

File tree

3 files changed

+65
-2
lines changed

3 files changed

+65
-2
lines changed
Lines changed: 53 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,53 @@
1+
# Narrowing with assert statements
2+
3+
## `assert` a value `is None` or `is not None`
4+
5+
```py
6+
def _(x: str | None, y: str | None):
7+
assert x is not None
8+
reveal_type(x) # revealed: str
9+
assert y is None
10+
reveal_type(y) # revealed: None
11+
```
12+
13+
## `assert` a value is truthy or falsy
14+
15+
```py
16+
def _(x: bool, y: bool):
17+
assert x
18+
reveal_type(x) # revealed: Literal[True]
19+
assert not y
20+
reveal_type(y) # revealed: Literal[False]
21+
```
22+
23+
## `assert` with `is` and `==` for literals
24+
25+
```py
26+
from typing import Literal
27+
28+
def _(x: Literal[1, 2, 3], y: Literal[1, 2, 3]):
29+
assert x is 2
30+
reveal_type(x) # revealed: Literal[2]
31+
assert y == 2
32+
reveal_type(y) # revealed: Literal[1, 2, 3]
33+
```
34+
35+
## `assert` with `isinstance`
36+
37+
```py
38+
def _(x: int | str):
39+
assert isinstance(x, int)
40+
reveal_type(x) # revealed: int
41+
```
42+
43+
## `assert` a value `in` a tuple
44+
45+
```py
46+
from typing import Literal
47+
48+
def _(x: Literal[1, 2, 3], y: Literal[1, 2, 3]):
49+
assert x in (1, 2)
50+
reveal_type(x) # revealed: Literal[1, 2]
51+
assert y not in (1, 2)
52+
reveal_type(y) # revealed: Literal[3]
53+
```

crates/red_knot_python_semantic/src/semantic_index/builder.rs

Lines changed: 11 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -535,7 +535,6 @@ impl<'db> SemanticIndexBuilder<'db> {
535535
}
536536

537537
/// Records a visibility constraint by applying it to all live bindings and declarations.
538-
#[must_use = "A visibility constraint must always be negated after it is added"]
539538
fn record_visibility_constraint(
540539
&mut self,
541540
predicate: Predicate<'db>,
@@ -1300,6 +1299,17 @@ where
13001299
);
13011300
}
13021301
}
1302+
1303+
ast::Stmt::Assert(node) => {
1304+
self.visit_expr(&node.test);
1305+
let predicate = self.record_expression_narrowing_constraint(&node.test);
1306+
self.record_visibility_constraint(predicate);
1307+
1308+
if let Some(msg) = &node.msg {
1309+
self.visit_expr(msg);
1310+
}
1311+
}
1312+
13031313
ast::Stmt::Assign(node) => {
13041314
debug_assert_eq!(&self.current_assignments, &[]);
13051315

crates/red_knot_python_semantic/src/types/infer.rs

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -3249,7 +3249,7 @@ impl<'db> TypeInferenceBuilder<'db> {
32493249
msg,
32503250
} = assert;
32513251

3252-
let test_ty = self.infer_expression(test);
3252+
let test_ty = self.infer_standalone_expression(test);
32533253

32543254
if let Err(err) = test_ty.try_bool(self.db()) {
32553255
err.report_diagnostic(&self.context, &**test);

0 commit comments

Comments
 (0)