Skip to content

Commit

Permalink
red-knot: support narrowing for bool(E) (#14668)
Browse files Browse the repository at this point in the history
Resolves #14547 by delegating
narrowing to `E` for `bool(E)` where `E` is some expression.

This change does not include other builtin class constructors which
should also work in this position, like `int(..)` or `float(..)`, as the
original issue does not mention these. It should be easy enough to add
checks for these as well if we want to.

I don't see a lot of markdown tests for malformed input, maybe there's a
better place for the no args and too many args cases to go?

I did see after the fact that it looks like this task was intended for a
new hire.. my apologies. I got here from
#13694, which is marked
help-wanted.

---------

Co-authored-by: David Peter <[email protected]>
  • Loading branch information
connorskees and sharkdp authored Dec 3, 2024
1 parent 91e2d9a commit 3e702e1
Show file tree
Hide file tree
Showing 2 changed files with 78 additions and 34 deletions.
Original file line number Diff line number Diff line change
@@ -0,0 +1,32 @@
## Narrowing for `bool(..)` checks

```py
def flag() -> bool: ...

x = 1 if flag() else None

# valid invocation, positive
reveal_type(x) # revealed: Literal[1] | None
if bool(x is not None):
reveal_type(x) # revealed: Literal[1]

# valid invocation, negative
reveal_type(x) # revealed: Literal[1] | None
if not bool(x is not None):
reveal_type(x) # revealed: None

# no args/narrowing
reveal_type(x) # revealed: Literal[1] | None
if not bool():
reveal_type(x) # revealed: Literal[1] | None

# invalid invocation, too many positional args
reveal_type(x) # revealed: Literal[1] | None
if bool(x is not None, 5): # TODO diagnostic
reveal_type(x) # revealed: Literal[1] | None

# invalid invocation, too many kwargs
reveal_type(x) # revealed: Literal[1] | None
if bool(x is not None, y=5): # TODO diagnostic
reveal_type(x) # revealed: Literal[1] | None
```
80 changes: 46 additions & 34 deletions crates/red_knot_python_semantic/src/types/narrow.rs
Original file line number Diff line number Diff line change
Expand Up @@ -388,46 +388,58 @@ impl<'db> NarrowingConstraintsBuilder<'db> {
let scope = self.scope();
let inference = infer_expression_types(self.db, expression);

let callable_ty =
inference.expression_ty(expr_call.func.scoped_expression_id(self.db, scope));

// TODO: add support for PEP 604 union types on the right hand side of `isinstance`
// and `issubclass`, for example `isinstance(x, str | (int | float))`.
match inference
.expression_ty(expr_call.func.scoped_expression_id(self.db, scope))
.into_function_literal()
.and_then(|f| f.known(self.db))
.and_then(KnownFunction::constraint_function)
{
Some(function) if expr_call.arguments.keywords.is_empty() => {
if let [ast::Expr::Name(ast::ExprName { id, .. }), class_info] =
match callable_ty {
Type::FunctionLiteral(function_type) if expr_call.arguments.keywords.is_empty() => {
let function = function_type
.known(self.db)
.and_then(KnownFunction::constraint_function)?;

let [ast::Expr::Name(ast::ExprName { id, .. }), class_info] =
&*expr_call.arguments.args
{
let symbol = self.symbols().symbol_id_by_name(id).unwrap();
else {
return None;
};

let class_info_ty =
inference.expression_ty(class_info.scoped_expression_id(self.db, scope));
let symbol = self.symbols().symbol_id_by_name(id).unwrap();

let to_constraint = match function {
KnownConstraintFunction::IsInstance => {
|class_literal: ClassLiteralType<'db>| {
Type::instance(class_literal.class)
}
}
KnownConstraintFunction::IsSubclass => {
|class_literal: ClassLiteralType<'db>| {
Type::subclass_of(class_literal.class)
}
}
};
let class_info_ty =
inference.expression_ty(class_info.scoped_expression_id(self.db, scope));

generate_classinfo_constraint(self.db, &class_info_ty, to_constraint).map(
|constraint| {
let mut constraints = NarrowingConstraints::default();
constraints.insert(symbol, constraint.negate_if(self.db, !is_positive));
constraints
},
)
} else {
None
}
let to_constraint = match function {
KnownConstraintFunction::IsInstance => {
|class_literal: ClassLiteralType<'db>| Type::instance(class_literal.class)
}
KnownConstraintFunction::IsSubclass => {
|class_literal: ClassLiteralType<'db>| {
Type::subclass_of(class_literal.class)
}
}
};

generate_classinfo_constraint(self.db, &class_info_ty, to_constraint).map(
|constraint| {
let mut constraints = NarrowingConstraints::default();
constraints.insert(symbol, constraint.negate_if(self.db, !is_positive));
constraints
},
)
}
// for the expression `bool(E)`, we further narrow the type based on `E`
Type::ClassLiteral(class_type)
if expr_call.arguments.args.len() == 1
&& expr_call.arguments.keywords.is_empty()
&& class_type.class.is_known(self.db, KnownClass::Bool) =>
{
self.evaluate_expression_node_constraint(
&expr_call.arguments.args[0],
expression,
is_positive,
)
}
_ => None,
}
Expand Down

0 comments on commit 3e702e1

Please sign in to comment.