Skip to content

Commit

Permalink
Recognize named type guards
Browse files Browse the repository at this point in the history
  • Loading branch information
InSyncWithFoo committed Feb 26, 2025
1 parent 89a76c7 commit ade4c79
Show file tree
Hide file tree
Showing 2 changed files with 41 additions and 3 deletions.
8 changes: 6 additions & 2 deletions crates/red_knot_python_semantic/src/types.rs
Original file line number Diff line number Diff line change
Expand Up @@ -3583,6 +3583,7 @@ pub enum KnownFunction {
Repr,
/// `typing(_extensions).final`
Final,
// TODO: Move this to `KnownClass`
/// `builtins.staticmethod`
StaticMethod,

Expand Down Expand Up @@ -3642,7 +3643,9 @@ impl KnownFunction {
/// Return `true` if `self` is defined in `module` at runtime.
const fn check_module(self, module: KnownModule) -> bool {
match self {
Self::IsInstance | Self::IsSubclass | Self::Len | Self::Repr | Self::StaticMethod => module.is_builtins(),
Self::IsInstance | Self::IsSubclass | Self::Len | Self::Repr | Self::StaticMethod => {
module.is_builtins()
}
Self::AssertType
| Self::Cast
| Self::Overload
Expand Down Expand Up @@ -4562,7 +4565,8 @@ pub(crate) mod tests {
KnownFunction::Len
| KnownFunction::Repr
| KnownFunction::IsInstance
| KnownFunction::IsSubclass => KnownModule::Builtins,
| KnownFunction::IsSubclass
| KnownFunction::StaticMethod => KnownModule::Builtins,

KnownFunction::GetattrStatic => KnownModule::Inspect,

Expand Down
36 changes: 35 additions & 1 deletion crates/red_knot_python_semantic/src/types/narrow.rs
Original file line number Diff line number Diff line change
Expand Up @@ -210,7 +210,7 @@ impl<'db> NarrowingConstraintsBuilder<'db> {
is_positive: bool,
) -> Option<NarrowingConstraints<'db>> {
match expression_node {
ast::Expr::Name(name) => Some(self.evaluate_expr_name(name, is_positive)),
ast::Expr::Name(name) => Some(self.evaluate_expr_name(name, expression, is_positive)),
ast::Expr::Compare(expr_compare) => {
self.evaluate_expr_compare(expr_compare, expression, is_positive)
}
Expand Down Expand Up @@ -257,16 +257,50 @@ impl<'db> NarrowingConstraintsBuilder<'db> {
fn evaluate_expr_name(
&mut self,
expr_name: &ast::ExprName,
expression: Expression<'db>,
is_positive: bool,
) -> NarrowingConstraints<'db> {
let ast::ExprName { id, .. } = expr_name;
let inference = infer_expression_types(self.db, expression);

let symbol = self
.symbols()
.symbol_id_by_name(id)
.expect("Should always have a symbol for every Name node");
let ty = inference.expression_type(expr_name.scoped_expression_id(self.db, self.scope()));

let mut constraints = NarrowingConstraints::default();

// TODO: Handle unions and intersections
let mut narrow_by_typeguards = || match ty {
Type::TypeGuard(type_guard) => {
let (_, guarded_symbol, _) = type_guard.symbol_info(self.db)?;

if !is_positive {
return None;
}

constraints.insert(
guarded_symbol,
type_guard.ty(self.db).negate_if(self.db, !is_positive),
);

Some(())
}
Type::TypeIs(type_is) => {
let (_, guarded_symbol, _) = type_is.symbol_info(self.db)?;

constraints.insert(
guarded_symbol,
type_is.ty(self.db).negate_if(self.db, !is_positive),
);

Some(())
}
_ => None,
};
narrow_by_typeguards();

constraints.insert(
symbol,
if is_positive {
Expand Down

0 comments on commit ade4c79

Please sign in to comment.