[red-knot] allow assignment expression in call compare narrowing (#17461)

## Summary

There was some narrowing constraints not covered from the previous PR

```py
def _(x: object):
    if (type(y := x)) is bool:
        reveal_type(y)  # revealed: bool
```

Also, refactored a bit

## Test Plan

Update type_api.md
This commit is contained in:
Matthew Mckee 2025-04-18 16:46:15 +01:00 committed by GitHub
parent 84d064a14c
commit 5853eb28dd
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
2 changed files with 32 additions and 29 deletions

View file

@ -144,3 +144,13 @@ def _(x: Base):
# express a constraint like `Base & ~ProperSubtypeOf[Base]`.
reveal_type(x) # revealed: Base
```
## Assignment expressions
```py
def _(x: object):
if (y := type(x)) is bool:
reveal_type(y) # revealed: Literal[bool]
if (type(y := x)) is bool:
reveal_type(y) # revealed: bool
```

View file

@ -238,6 +238,17 @@ fn negate_if<'db>(constraints: &mut NarrowingConstraints<'db>, db: &'db dyn Db,
}
}
fn expr_name(expr: &ast::Expr) -> Option<&ast::name::Name> {
match expr {
ast::Expr::Named(ast::ExprNamed { target, .. }) => match target.as_ref() {
ast::Expr::Name(ast::ExprName { id, .. }) => Some(id),
_ => None,
},
ast::Expr::Name(ast::ExprName { id, .. }) => Some(id),
_ => None,
}
}
struct NarrowingConstraintsBuilder<'db> {
db: &'db dyn Db,
predicate: PredicateNode<'db>,
@ -497,27 +508,9 @@ impl<'db> NarrowingConstraintsBuilder<'db> {
last_rhs_ty = Some(rhs_ty);
match left {
ast::Expr::Name(ast::ExprName {
range: _,
id,
ctx: _,
}) => {
let symbol = self.expect_expr_name_symbol(id);
let op = if is_positive { *op } else { op.negate() };
if let Some(ty) = self.evaluate_expr_compare_op(lhs_ty, rhs_ty, op) {
constraints.insert(symbol, ty);
}
}
ast::Expr::Named(ast::ExprNamed {
range: _,
target,
value: _,
}) => {
if let ast::Expr::Name(ast::ExprName { id, .. }) = target.as_ref() {
ast::Expr::Name(_) | ast::Expr::Named(_) => {
if let Some(id) = expr_name(left) {
let symbol = self.expect_expr_name_symbol(id);
let op = if is_positive { *op } else { op.negate() };
if let Some(ty) = self.evaluate_expr_compare_op(lhs_ty, rhs_ty, op) {
@ -545,8 +538,12 @@ impl<'db> NarrowingConstraintsBuilder<'db> {
}
};
let [ast::Expr::Name(ast::ExprName { id, .. })] = &**args else {
continue;
let id = match &**args {
[first] => match expr_name(first) {
Some(id) => id,
None => continue,
},
_ => continue,
};
let is_valid_constraint = if is_positive {
@ -598,13 +595,9 @@ impl<'db> NarrowingConstraintsBuilder<'db> {
let function = function_type.known(self.db)?.into_constraint_function()?;
let (id, class_info) = match &*expr_call.arguments.args {
[first, class_info] => match first {
ast::Expr::Named(ast::ExprNamed { target, .. }) => match target.as_ref() {
ast::Expr::Name(ast::ExprName { id, .. }) => (id, class_info),
_ => return None,
},
ast::Expr::Name(ast::ExprName { id, .. }) => (id, class_info),
_ => return None,
[first, class_info] => match expr_name(first) {
Some(id) => (id, class_info),
None => return None,
},
_ => return None,
};