[red-knot] Implement more types in binary and unary expressions (#13803)

Implemented some points from
#12701

- Handle Unknown and Any in Unary operation
- Handle Boolean in binary operations
- Handle instances in unary operation
- Consider division by False to be division by zero

---------

Co-authored-by: Carl Meyer <[email protected]>
Co-authored-by: Alex Waygood <[email protected]>

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

@@ -11,6 +11,18 @@ reveal_type(-3 / 3)  # revealed: float
 reveal_type(5 % 3)  # revealed: Literal[2]
 ```
 
+## Power
+
+For power if the result fits in the int literal type it will be a Literal type. Otherwise the
+outcome is int.
+
+```py
+largest_u32 = 4_294_967_295
+reveal_type(2**2)  # revealed: Literal[4]
+reveal_type(1 ** (largest_u32 + 1))  # revealed: int
+reveal_type(2**largest_u32)  # revealed: int
+```
+
 ## Division by Zero
 
 This error is really outside the current Python type system, because e.g. `int.__truediv__` and
@@ -38,6 +50,14 @@ reveal_type(c)  # revealed: int
 # revealed: float
 reveal_type(int() / 0)
 
+# error: "Cannot divide object of type `Literal[1]` by zero"
+# revealed: float
+reveal_type(1 / False)
+# error: [division-by-zero] "Cannot divide object of type `Literal[True]` by zero"
+True / False
+# error: [division-by-zero] "Cannot divide object of type `Literal[True]` by zero"
+bool(1) / False
+
 # error: "Cannot divide object of type `float` by zero"
 # revealed: float
 reveal_type(1.0 / 0)

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

@@ -0,0 +1,32 @@
+# 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] "Unary operator `+` is unsupported for type `NoDunder`"
+-b  # error: [unsupported-operator] "Unary operator `-` is unsupported for type `NoDunder`"
+~b  # error: [unsupported-operator] "Unary operator `~` is unsupported for type `NoDunder`"
+```

crates/red_knot_python_semantic/src/types/infer.rs

@@ -539,10 +539,11 @@ impl<'db> TypeInferenceBuilder<'db> {
     /// 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>) {
         match left {
-            Type::IntLiteral(_) => {}
+            Type::BooleanLiteral(_) | 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,
         };
 
@@ -2459,7 +2460,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 +2472,35 @@ 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 unary_dunder_method = match op {
+                    UnaryOp::Invert => "__invert__",
+                    UnaryOp::UAdd => "__pos__",
+                    UnaryOp::USub => "__neg__",
+                    UnaryOp::Not => {
+                        unreachable!("Not operator is handled in its own case");
+                    }
+                };
+                let class_member = class.class_member(self.db, unary_dunder_method);
+                let call = class_member.call(self.db, &[operand_type]);
 
+                match call.return_ty_result(self.db, AnyNodeRef::ExprUnaryOp(unary), self) {
+                    Ok(t) => t,
+                    Err(e) => {
+                        self.add_diagnostic(
+                            unary.into(),
+                            "unsupported-operator",
+                            format_args!(
+                                "Unary operator `{op}` is unsupported for type `{}`",
+                                operand_type.display(self.db),
+                            ),
+                        );
+                        e.return_ty()
+                    }
+                }
+            }
             _ => Type::Todo, // TODO other unary op types
         }
     }
@@ -2491,7 +2522,7 @@ impl<'db> TypeInferenceBuilder<'db> {
             (op, right_ty),
             (
                 ast::Operator::Div | ast::Operator::FloorDiv | ast::Operator::Mod,
-                Type::IntLiteral(0),
+                Type::IntLiteral(0) | Type::BooleanLiteral(false)
             )
         ) {
             self.check_division_by_zero(binary, left_ty);
@@ -2558,6 +2589,17 @@ impl<'db> TypeInferenceBuilder<'db> {
                     .unwrap_or_else(|| KnownClass::Int.to_instance(self.db)),
             ),
 
+            (Type::IntLiteral(n), Type::IntLiteral(m), ast::Operator::Pow) => {
+                let m = u32::try_from(m);
+                Some(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) => {
                 Some(Type::BytesLiteral(BytesLiteralType::new(
                     self.db,
@@ -2693,6 +2735,20 @@ impl<'db> TypeInferenceBuilder<'db> {
                     })
             }
 
+            (
+                Type::BooleanLiteral(b1),
+                Type::BooleanLiteral(b2),
+                ruff_python_ast::Operator::BitOr,
+            ) => Some(Type::BooleanLiteral(b1 | b2)),
+
+            (Type::BooleanLiteral(bool_value), right, op) => self.infer_binary_expression_type(
+                Type::IntLiteral(i64::from(bool_value)),
+                right,
+                op,
+            ),
+            (left, Type::BooleanLiteral(bool_value), op) => {
+                self.infer_binary_expression_type(left, Type::IntLiteral(i64::from(bool_value)), op)
+            }
             _ => Some(Type::Todo), // TODO
         }
     }