mirror of
https://github.com/astral-sh/ruff.git
synced 2025-09-30 22:01:18 +00:00
[red-knot] arithmetic on int literals (#11760)
## Summary Add support for inferring int literal types from basic arithmetic on int literals. Just to begin showing examples of resolving more complex expression types, and because this will be useful in testing walrus expressions. ## Test Plan Added test.
This commit is contained in:
parent
9b2cf569b2
commit
b46e9e825a
2 changed files with 85 additions and 3 deletions
|
@ -8,6 +8,7 @@ use crate::semantic::{
|
||||||
};
|
};
|
||||||
use crate::{FxDashMap, FxIndexSet, Name};
|
use crate::{FxDashMap, FxIndexSet, Name};
|
||||||
use ruff_index::{newtype_index, IndexVec};
|
use ruff_index::{newtype_index, IndexVec};
|
||||||
|
use ruff_python_ast as ast;
|
||||||
use rustc_hash::FxHashMap;
|
use rustc_hash::FxHashMap;
|
||||||
|
|
||||||
pub(crate) mod infer;
|
pub(crate) mod infer;
|
||||||
|
@ -55,9 +56,9 @@ impl Type {
|
||||||
|
|
||||||
pub fn get_member(&self, db: &dyn SemanticDb, name: &Name) -> QueryResult<Option<Type>> {
|
pub fn get_member(&self, db: &dyn SemanticDb, name: &Name) -> QueryResult<Option<Type>> {
|
||||||
match self {
|
match self {
|
||||||
Type::Any => todo!("attribute lookup on Any type"),
|
Type::Any => Ok(Some(Type::Any)),
|
||||||
Type::Never => todo!("attribute lookup on Never type"),
|
Type::Never => todo!("attribute lookup on Never type"),
|
||||||
Type::Unknown => todo!("attribute lookup on Unknown type"),
|
Type::Unknown => Ok(Some(Type::Unknown)),
|
||||||
Type::Unbound => todo!("attribute lookup on Unbound type"),
|
Type::Unbound => todo!("attribute lookup on Unbound type"),
|
||||||
Type::Function(_) => todo!("attribute lookup on Function type"),
|
Type::Function(_) => todo!("attribute lookup on Function type"),
|
||||||
Type::Module(module_id) => module_id.get_member(db, name),
|
Type::Module(module_id) => module_id.get_member(db, name),
|
||||||
|
@ -85,6 +86,56 @@ impl Type {
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
|
// when this is fully fleshed out, it will use the db arg and may return QueryError
|
||||||
|
#[allow(clippy::unnecessary_wraps)]
|
||||||
|
pub fn resolve_bin_op(
|
||||||
|
&self,
|
||||||
|
_db: &dyn SemanticDb,
|
||||||
|
op: ast::Operator,
|
||||||
|
right_ty: Type,
|
||||||
|
) -> QueryResult<Type> {
|
||||||
|
match self {
|
||||||
|
Type::Any => Ok(Type::Any),
|
||||||
|
Type::Unknown => Ok(Type::Unknown),
|
||||||
|
Type::IntLiteral(n) => {
|
||||||
|
match right_ty {
|
||||||
|
Type::IntLiteral(m) => {
|
||||||
|
match op {
|
||||||
|
ast::Operator::Add => Ok(n
|
||||||
|
.checked_add(m)
|
||||||
|
.map(Type::IntLiteral)
|
||||||
|
// TODO builtins.int
|
||||||
|
.unwrap_or(Type::Unknown)),
|
||||||
|
ast::Operator::Sub => Ok(n
|
||||||
|
.checked_sub(m)
|
||||||
|
.map(Type::IntLiteral)
|
||||||
|
// TODO builtins.int
|
||||||
|
.unwrap_or(Type::Unknown)),
|
||||||
|
ast::Operator::Mult => Ok(n
|
||||||
|
.checked_mul(m)
|
||||||
|
.map(Type::IntLiteral)
|
||||||
|
// TODO builtins.int
|
||||||
|
.unwrap_or(Type::Unknown)),
|
||||||
|
ast::Operator::Div => Ok(n
|
||||||
|
.checked_div(m)
|
||||||
|
.map(Type::IntLiteral)
|
||||||
|
// TODO builtins.int
|
||||||
|
.unwrap_or(Type::Unknown)),
|
||||||
|
ast::Operator::Mod => Ok(n
|
||||||
|
.checked_rem(m)
|
||||||
|
.map(Type::IntLiteral)
|
||||||
|
// TODO division by zero error
|
||||||
|
.unwrap_or(Type::Unknown)),
|
||||||
|
_ => todo!("complete binop op support for IntLiteral"),
|
||||||
|
}
|
||||||
|
}
|
||||||
|
_ => todo!("complete binop right_ty support for IntLiteral"),
|
||||||
|
}
|
||||||
|
}
|
||||||
|
_ => todo!("complete binop support"),
|
||||||
|
}
|
||||||
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
impl From<FunctionTypeId> for Type {
|
impl From<FunctionTypeId> for Type {
|
||||||
|
|
|
@ -212,7 +212,15 @@ fn infer_expr_type(db: &dyn SemanticDb, file_id: FileId, expr: &ast::Expr) -> Qu
|
||||||
.get_member(db, attr_name)
|
.get_member(db, attr_name)
|
||||||
.map(|ty| ty.unwrap_or(Type::Unknown))
|
.map(|ty| ty.unwrap_or(Type::Unknown))
|
||||||
}
|
}
|
||||||
_ => todo!("full expression type resolution"),
|
ast::Expr::BinOp(ast::ExprBinOp {
|
||||||
|
left, op, right, ..
|
||||||
|
}) => {
|
||||||
|
let left_ty = infer_expr_type(db, file_id, left)?;
|
||||||
|
let right_ty = infer_expr_type(db, file_id, right)?;
|
||||||
|
// TODO add reverse bin op support if right <: left
|
||||||
|
left_ty.resolve_bin_op(db, *op, right_ty)
|
||||||
|
}
|
||||||
|
_ => todo!("expression type resolution for {:?}", expr),
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
|
@ -488,4 +496,27 @@ mod tests {
|
||||||
|
|
||||||
assert_public_type(&case, "a", "x", "(Literal[2] | Literal[3] | Literal[4])")
|
assert_public_type(&case, "a", "x", "(Literal[2] | Literal[3] | Literal[4])")
|
||||||
}
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn literal_int_arithmetic() -> anyhow::Result<()> {
|
||||||
|
let case = create_test()?;
|
||||||
|
|
||||||
|
write_to_path(
|
||||||
|
&case,
|
||||||
|
"a.py",
|
||||||
|
"
|
||||||
|
a = 2 + 1
|
||||||
|
b = a - 4
|
||||||
|
c = a * b
|
||||||
|
d = c / 3
|
||||||
|
e = 5 % 3
|
||||||
|
",
|
||||||
|
)?;
|
||||||
|
|
||||||
|
assert_public_type(&case, "a", "a", "Literal[3]")?;
|
||||||
|
assert_public_type(&case, "a", "b", "Literal[-1]")?;
|
||||||
|
assert_public_type(&case, "a", "c", "Literal[-3]")?;
|
||||||
|
assert_public_type(&case, "a", "d", "Literal[-1]")?;
|
||||||
|
assert_public_type(&case, "a", "e", "Literal[2]")
|
||||||
|
}
|
||||||
}
|
}
|
||||||
|
|
Loading…
Add table
Add a link
Reference in a new issue