diff --git a/crates/red_knot/src/semantic/types.rs b/crates/red_knot/src/semantic/types.rs index 24178165cc..8db0d7da9c 100644 --- a/crates/red_knot/src/semantic/types.rs +++ b/crates/red_knot/src/semantic/types.rs @@ -8,6 +8,7 @@ use crate::semantic::{ }; use crate::{FxDashMap, FxIndexSet, Name}; use ruff_index::{newtype_index, IndexVec}; +use ruff_python_ast as ast; use rustc_hash::FxHashMap; pub(crate) mod infer; @@ -55,9 +56,9 @@ impl Type { pub fn get_member(&self, db: &dyn SemanticDb, name: &Name) -> QueryResult> { 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::Unknown => todo!("attribute lookup on Unknown type"), + Type::Unknown => Ok(Some(Type::Unknown)), Type::Unbound => todo!("attribute lookup on Unbound type"), Type::Function(_) => todo!("attribute lookup on Function type"), 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 { + 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 for Type { diff --git a/crates/red_knot/src/semantic/types/infer.rs b/crates/red_knot/src/semantic/types/infer.rs index 20d25c06db..7e8c9c4139 100644 --- a/crates/red_knot/src/semantic/types/infer.rs +++ b/crates/red_knot/src/semantic/types/infer.rs @@ -212,7 +212,15 @@ fn infer_expr_type(db: &dyn SemanticDb, file_id: FileId, expr: &ast::Expr) -> Qu .get_member(db, attr_name) .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])") } + + #[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]") + } }