[red-knot] Implement more types in binary and unary expressions

crates/red_knot_python_semantic/resources/mdtest/binary/booleans.md

@@ -0,0 +1,48 @@
+## Binary operations on booleans
+
+## Basic Arithmetic
+
+We try to be precise and all operations except for division will result in Literal type.
+
+```py
+a = True
+b = False
+
+reveal_type(a + a)  # revealed: Literal[2]
+reveal_type(a + b)  # revealed: Literal[1]
+reveal_type(b + a)  # revealed: Literal[1]
+reveal_type(b + b)  # revealed: Literal[0]
+
+reveal_type(a - a)  # revealed: Literal[0]
+reveal_type(a - b)  # revealed: Literal[1]
+reveal_type(b - a)  # revealed: Literal[-1]
+reveal_type(b - b)  # revealed: Literal[0]
+
+reveal_type(a * a)  # revealed: Literal[1]
+reveal_type(a * b)  # revealed: Literal[0]
+reveal_type(b * a)  # revealed: Literal[0]
+reveal_type(b * b)  # revealed: Literal[0]
+
+reveal_type(a % a)  # revealed: Literal[0]
+reveal_type(b % a)  # revealed: Literal[0]
+
+reveal_type(a // a)  # revealed: Literal[1]
+reveal_type(b // a)  # revealed: Literal[0]
+
+reveal_type(a ** a)  # revealed: Literal[1]
+reveal_type(a ** b)  # revealed: Literal[1]
+reveal_type(b ** a)  # revealed: Literal[0]
+reveal_type(b ** b)  # revealed: Literal[1]
+
+# Division
+reveal_type(a / a)  # revealed: float
+reveal_type(b / a)  # revealed: float
+b / b  # error: [division-by-zero] "Cannot divide object of type `Literal[False]` by zero"
+a / b  # error: [division-by-zero] "Cannot divide object of type `Literal[True]` by zero"
+
+# bitwise OR
+reveal_type(a | a)  # revealed: Literal[True]
+reveal_type(a | b)  # revealed: Literal[True]
+reveal_type(b | a)  # revealed: Literal[True]
+reveal_type(b | b)  # revealed: Literal[False]
+```

crates/red_knot_python_semantic/resources/mdtest/binary/integers.md

@@ -26,6 +26,7 @@ subclass; we only emit the error if the LHS type is exactly `int` or `float`, no
 
 ```py
 a = 1 / 0  # error: "Cannot divide object of type `Literal[1]` by zero"
+
 reveal_type(a)  # revealed: float
 
 b = 2 // 0  # error: "Cannot floor divide object of type `Literal[2]` by zero"
@@ -38,6 +39,11 @@ d = int() / 0  # error: "Cannot divide object of type `int` by zero"
 # TODO should be int
 reveal_type(d)  # revealed: @Todo
 
+f = 1 / False # error: "Cannot divide object of type `Literal[1]` by zero"
+reveal_type(f)  # revealed: float
+True / False # error: [division-by-zero] "Cannot divide object of type `Literal[True]` by zero"
+bool(1) / False # error: [division-by-zero] "Cannot divide object of type `Literal[True]` by zero"
+
 e = 1.0 / 0  # error: "Cannot divide object of type `float` by zero"
 # TODO should be float
 reveal_type(e)  # revealed: @Todo

crates/red_knot_python_semantic/resources/mdtest/unary/instance.md

@@ -0,0 +1,31 @@
+# Unary Operations
+
+```py
+class Number:
+    def __init__(self, value: int):
+        self.value = 1
+
+    def __pos__(self) -> int:
+        return +self.value
+
+    def __neg__(self) -> int:
+        return -self.value
+
+    def __invert__(self) -> Literal[True]:
+        return True
+
+a = Number()
+
+reveal_type(+a) # revealed: int
+reveal_type(-a) # revealed: int
+reveal_type(~a) # revealed: @Todo
+
+class NoDunder:
+  ...
+
+b = NoDunder()
++b # error: [unsupported-operator] "Operator `+` is unsupported for type `NoDunder`"
+-b # error: [unsupported-operator] "Operator `-` is unsupported for type `NoDunder`"
+~b # error: [unsupported-operator] "Operator `~` is unsupported for type `NoDunder`"
+
+```

crates/red_knot_python_semantic/src/types.rs

@@ -355,6 +355,13 @@ impl<'db> Type<'db> {
     pub const fn into_int_literal_type(self) -> Option<i64> {
         match self {
             Type::IntLiteral(value) => Some(value),
+            Type::BooleanLiteral(b) => {
+                if b {
+                    Some(1)
+                } else {
+                    Some(0)
+                }
+            }
             _ => None,
         }
     }

crates/red_knot_python_semantic/src/types/infer.rs

@@ -57,7 +57,7 @@ use crate::types::{
 };
 use crate::Db;
 
-use super::{KnownClass, UnionBuilder};
+use super::{CallOutcome, KnownClass, UnionBuilder};
 
 /// Infer all types for a [`ScopeId`], including all definitions and expressions in that scope.
 /// Use when checking a scope, or needing to provide a type for an arbitrary expression in the
@@ -537,12 +537,18 @@ impl<'db> TypeInferenceBuilder<'db> {
     /// Raise a diagnostic if the given type cannot be divided by zero.
     ///
     /// Expects the resolved type of the left side of the binary expression.
-    fn check_division_by_zero(&mut self, expr: &ast::ExprBinOp, left: Type<'db>) {
+    fn check_division_by_zero(
+        &mut self,
+        expr: &ast::ExprBinOp,
+        left: Type<'db>,
+        original_left: Type<'db>,
+    ) {
         match left {
             Type::IntLiteral(_) => {}
             Type::Instance(cls)
-                if cls.is_known(self.db, KnownClass::Float)
-                    || cls.is_known(self.db, KnownClass::Int) => {}
+                if [KnownClass::Float, KnownClass::Int, KnownClass::Bool]
+                    .iter()
+                    .any(|&k| cls.is_known(self.db, k)) => {}
             _ => return,
         };
 
@@ -558,7 +564,7 @@ impl<'db> TypeInferenceBuilder<'db> {
             "division-by-zero",
             format_args!(
                 "Cannot {op} object of type `{}` {by_zero}",
-                left.display(self.db)
+                original_left.display(self.db)
             ),
         );
     }
@@ -2459,7 +2465,9 @@ impl<'db> TypeInferenceBuilder<'db> {
             operand,
         } = unary;
 
-        match (op, self.infer_expression(operand)) {
+        let operand_type = self.infer_expression(operand);
+
+        match (op, operand_type) {
             (UnaryOp::UAdd, Type::IntLiteral(value)) => Type::IntLiteral(value),
             (UnaryOp::USub, Type::IntLiteral(value)) => Type::IntLiteral(-value),
             (UnaryOp::Invert, Type::IntLiteral(value)) => Type::IntLiteral(!value),
@@ -2469,7 +2477,24 @@ impl<'db> TypeInferenceBuilder<'db> {
             (UnaryOp::Invert, Type::BooleanLiteral(bool)) => Type::IntLiteral(!i64::from(bool)),
 
             (UnaryOp::Not, ty) => ty.bool(self.db).negate().into_type(self.db),
+            (_, Type::Any) => Type::Any,
+            (_, Type::Unknown) => Type::Unknown,
+            (op @ (UnaryOp::UAdd | UnaryOp::USub | UnaryOp::Invert), Type::Instance(class)) => {
+                let class_member = class.class_member(self.db, op.dunder());
+                let call = class_member.call(self.db, &[operand_type]);
+                if matches!(call, CallOutcome::NotCallable { not_callable_ty: _ }) {
+                    self.add_diagnostic(
+                        unary.into(),
+                        "unsupported-operator",
+                        format_args!(
+                            "Operator `{op}` is unsupported for type `{}`",
+                            operand_type.display(self.db),
+                        ),
+                    );
+                }
 
+                call.return_ty(self.db).unwrap_or(Type::Unknown)
+            }
             _ => Type::Todo, // TODO other unary op types
         }
     }
@@ -2484,17 +2509,32 @@ impl<'db> TypeInferenceBuilder<'db> {
 
         let left_ty = self.infer_expression(left);
         let right_ty = self.infer_expression(right);
+        self.infer_binary_expression_type(binary, *op, left_ty, right_ty)
+    }
 
+    fn infer_binary_expression_type(
+        &mut self,
+        expr: &ast::ExprBinOp,
+        op: ast::Operator,
+        left_ty: Type<'db>,
+        right_ty: Type<'db>,
+    ) -> Type<'db> {
         // Check for division by zero; this doesn't change the inferred type for the expression, but
         // may emit a diagnostic
         if matches!(
             (op, right_ty),
             (
                 ast::Operator::Div | ast::Operator::FloorDiv | ast::Operator::Mod,
-                Type::IntLiteral(0),
+                Type::IntLiteral(0)
             )
         ) {
-            self.check_division_by_zero(binary, left_ty);
+            // We want to print the original type that was used in the expression. The left type might
+            // have been converted to other types during the check.
+            // Since the type was previously inferred in the function that calls this or in the
+            // previous recursion this id will exist
+            let expr_id = expr.left.scoped_ast_id(self.db, self.scope());
+            let previous = self.types.expressions[&expr_id];
+            self.check_division_by_zero(expr, left_ty, previous);
         }
 
         match (left_ty, right_ty, op) {
@@ -2532,6 +2572,17 @@ impl<'db> TypeInferenceBuilder<'db> {
                 .map(Type::IntLiteral)
                 .unwrap_or_else(|| KnownClass::Int.to_instance(self.db)),
 
+            (Type::IntLiteral(n), Type::IntLiteral(m), ast::Operator::Pow) => {
+                let m = u32::try_from(m);
+                match m {
+                    Ok(m) => n
+                        .checked_pow(m)
+                        .map(Type::IntLiteral)
+                        .unwrap_or_else(|| KnownClass::Int.to_instance(self.db)),
+                    Err(_) => KnownClass::Int.to_instance(self.db),
+                }
+            }
+
             (Type::BytesLiteral(lhs), Type::BytesLiteral(rhs), ast::Operator::Add) => {
                 Type::BytesLiteral(BytesLiteralType::new(
                     self.db,
@@ -2589,6 +2640,25 @@ impl<'db> TypeInferenceBuilder<'db> {
                 }
             }
 
+            (
+                Type::BooleanLiteral(b1),
+                Type::BooleanLiteral(b2),
+                ruff_python_ast::Operator::BitOr,
+            ) => match (b1, b2) {
+                (true, true) => Type::BooleanLiteral(true),
+                (true, false) => Type::BooleanLiteral(true),
+                (false, true) => Type::BooleanLiteral(true),
+                (false, false) => Type::BooleanLiteral(false),
+            },
+
+            (Type::BooleanLiteral(_), right, op) => {
+                let int_value = left_ty.expect_int_literal();
+                self.infer_binary_expression_type(expr, op, Type::IntLiteral(int_value), right)
+            }
+            (left, Type::BooleanLiteral(_), op) => {
+                let int_value = right_ty.expect_int_literal();
+                self.infer_binary_expression_type(expr, op, left, Type::IntLiteral(int_value))
+            }
             _ => Type::Todo, // TODO
         }
     }

crates/ruff_python_ast/src/nodes.rs

@@ -2997,6 +2997,14 @@ impl UnaryOp {
             UnaryOp::USub => "-",
         }
     }
+    pub const fn dunder(self) -> &'static str {
+        match self {
+            UnaryOp::Invert => "__invert__",
+            UnaryOp::Not => "__bool__",
+            UnaryOp::UAdd => "__pos__",
+            UnaryOp::USub => "__neg__",
+        }
+    }
 }
 
 impl fmt::Display for UnaryOp {