[red-knot] Infer boolean literal expression (#12688)

## Summary

This PR implements type inference for boolean literal expressions.

## Test Plan

Add test cases for `True` and `False`.
This commit is contained in:
Dhruv Manilawala 2024-08-06 00:00:53 +05:30 committed by GitHub
parent 0b4d3ce39b
commit a8e2ba508e
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
3 changed files with 22 additions and 6 deletions

View file

@ -130,6 +130,8 @@ pub enum Type<'db> {
Union(UnionType<'db>), Union(UnionType<'db>),
Intersection(IntersectionType<'db>), Intersection(IntersectionType<'db>),
IntLiteral(i64), IntLiteral(i64),
/// A boolean literal, either `True` or `False`.
BooleanLiteral(bool),
// TODO protocols, callable types, overloads, generics, type vars // TODO protocols, callable types, overloads, generics, type vars
} }
@ -175,6 +177,7 @@ impl<'db> Type<'db> {
// TODO raise error // TODO raise error
Type::Unknown Type::Unknown
} }
Type::BooleanLiteral(_) => Type::Unknown,
} }
} }
} }

View file

@ -35,6 +35,9 @@ impl Display for DisplayType<'_> {
Type::Union(union) => union.display(self.db).fmt(f), Type::Union(union) => union.display(self.db).fmt(f),
Type::Intersection(intersection) => intersection.display(self.db).fmt(f), Type::Intersection(intersection) => intersection.display(self.db).fmt(f),
Type::IntLiteral(n) => write!(f, "Literal[{n}]"), Type::IntLiteral(n) => write!(f, "Literal[{n}]"),
Type::BooleanLiteral(boolean) => {
write!(f, "Literal[{}]", if *boolean { "True" } else { "False" })
}
} }
} }
} }

View file

@ -939,12 +939,10 @@ impl<'db> TypeInferenceBuilder<'db> {
} }
#[allow(clippy::unused_self)] #[allow(clippy::unused_self)]
fn infer_boolean_literal_expression( fn infer_boolean_literal_expression(&mut self, literal: &ast::ExprBooleanLiteral) -> Type<'db> {
&mut self, let ast::ExprBooleanLiteral { range: _, value } = literal;
_literal: &ast::ExprBooleanLiteral,
) -> Type<'db> { Type::BooleanLiteral(*value)
// TODO builtins.bool and boolean Literal types
Type::Unknown
} }
#[allow(clippy::unused_self)] #[allow(clippy::unused_self)]
@ -1649,6 +1647,18 @@ mod tests {
Ok(()) Ok(())
} }
#[test]
fn boolean_literal() -> anyhow::Result<()> {
let mut db = setup_db();
db.write_file("src/a.py", "x = True\ny = False")?;
assert_public_ty(&db, "src/a.py", "x", "Literal[True]");
assert_public_ty(&db, "src/a.py", "y", "Literal[False]");
Ok(())
}
#[test] #[test]
fn resolve_union() -> anyhow::Result<()> { fn resolve_union() -> anyhow::Result<()> {
let mut db = setup_db(); let mut db = setup_db();