[red-knot] WIP: binary arithmetic on instances

Co-authored-by: Alex Waygood <Alex.Waygood@Gmail.com>

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

@@ -0,0 +1,264 @@
+# Binary operations on instances
+
+Binary operations in Python are implemented by means of magic double-underscore methods.
+
+For references, see:
+
+- <https://snarky.ca/unravelling-binary-arithmetic-operations-in-python/>
+- <https://docs.python.org/3/reference/datamodel.html#emulating-numeric-types>
+
+## Operations
+
+We support inference for all Python's binary operators:
+`+`, `-`, `*`, `@`, `/`, `//`, `%`, `**`, `<<`, `>>`, `&`, `^`, and `|`.
+
+```py
+class A:
+    def __add__(self, other) -> A: return self
+    def __sub__(self, other) -> A: return self
+    def __mul__(self, other) -> A: return self
+    def __matmul__(self, other) -> A: return self
+    def __truediv__(self, other) -> A: return self
+    def __floordiv__(self, other) -> A: return self
+    def __mod__(self, other) -> A: return self
+    def __pow__(self, other) -> A: return self
+    def __lshift__(self, other) -> A: return self
+    def __rshift__(self, other) -> A: return self
+    def __and__(self, other) -> A: return self
+    def __xor__(self, other) -> A: return self
+    def __or__(self, other) -> A: return self
+
+class B: pass
+
+reveal_type(A() + B())  # revealed: A
+reveal_type(A() - B())  # revealed: A
+reveal_type(A() * B())  # revealed: A
+reveal_type(A() @ B())  # revealed: A
+reveal_type(A() / B())  # revealed: A
+reveal_type(A() // B())  # revealed: A
+reveal_type(A() % B())  # revealed: A
+reveal_type(A() ** B())  # revealed: A
+reveal_type(A() << B())  # revealed: A
+reveal_type(A() >> B())  # revealed: A
+reveal_type(A() & B())  # revealed: A
+reveal_type(A() ^ B())  # revealed: A
+reveal_type(A() | B())  # revealed: A
+```
+
+## Reflected
+
+We also support inference for reflected operations:
+
+```py
+class A:
+    def __radd__(self, other) -> A: return self
+    def __rsub__(self, other) -> A: return self
+    def __rmul__(self, other) -> A: return self
+    def __rmatmul__(self, other) -> A: return self
+    def __rtruediv__(self, other) -> A: return self
+    def __rfloordiv__(self, other) -> A: return self
+    def __rmod__(self, other) -> A: return self
+    def __rpow__(self, other) -> A: return self
+    def __rlshift__(self, other) -> A: return self
+    def __rrshift__(self, other) -> A: return self
+    def __rand__(self, other) -> A: return self
+    def __rxor__(self, other) -> A: return self
+    def __ror__(self, other) -> A: return self
+
+class B: pass
+
+
+reveal_type(B() + A())  # revealed: A
+reveal_type(B() - A())  # revealed: A
+reveal_type(B() * A())  # revealed: A
+reveal_type(B() @ A())  # revealed: A
+reveal_type(B() / A())  # revealed: A
+reveal_type(B() // A())  # revealed: A
+reveal_type(B() % A())  # revealed: A
+reveal_type(B() ** A())  # revealed: A
+reveal_type(B() << A())  # revealed: A
+reveal_type(B() >> A())  # revealed: A
+reveal_type(B() & A())  # revealed: A
+reveal_type(B() ^ A())  # revealed: A
+reveal_type(B() | A())  # revealed: A
+```
+
+## Return different type
+
+The magic methods aren't required to return the type of `self`:
+
+```py
+class A:
+    def __add__(self, other) -> int: return 1
+    def __rsub__(self, other) -> int: return 1
+
+reveal_type(A() + "foo")  # revealed: int
+reveal_type("foo" - A())  # revealed: int
+```
+
+## Reflected precedence
+
+If the right-hand operand is a subtype of the left-hand operand and has a
+different implementation of the reflected method, the reflected method on the
+right-hand operand takes precedence.
+
+```py
+class A:
+    def __add__(self, other) -> str: return "foo"
+
+class B(A):
+    def __radd__(self, other) -> int: return 42
+
+reveal_type(A() + B())  # revealed: int
+```
+
+## Reflected precedence 2
+
+If the right-hand operand is a subtype of the left-hand operand, but does not
+override the reflected method, the left-hand operand's non-reflected method
+still takes precedence:
+
+```py
+class A:
+    def __add__(self, other) -> str: return "foo"
+    def __radd__(self, other) -> int: return 42
+
+class B(A): pass
+
+reveal_type(A() + B())  # revealed: str
+```
+
+## Only reflected supported
+
+For example, at runtime, `(1).__add__(1.2)` is `NotImplemented`, but
+`(1.2).__radd__(1) == 2.2`, meaning that `1 + 1.2` succeeds at runtime
+(producing `2.2`).
+
+```py
+class A:
+    def __sub__(self, other: A) -> A:
+        return A()
+
+class B:
+    def __rsub__(self, other: A) -> B:
+        return B()
+
+A() - B()
+```
+
+## Callable instances as dunders
+
+Believe it or not, this is supported at runtime:
+
+```py
+class A:
+    def __call__(self, other) -> int:
+        return 42
+
+class B:
+    __add__ = A()
+
+reveal_type(B() + B())  # revealed: int
+```
+
+## Integration test: numbers from typeshed
+
+```py
+# TODO should be complex, need to check arg type and fall back
+reveal_type(3.14 + 3j)  # revealed: float
+reveal_type(3j + 3.14)  # revealed: complex
+# TODO should be float, need to check arg type and fall back
+reveal_type(42 + 4.2)  # revealed: int
+reveal_type(4.2 + 42)  # revealed: float
+reveal_type(3j + 3)  # revealed: complex
+# TODO should be complex, need to check arg type and fall back
+reveal_type(3 + 3j)  # revealed: int
+
+def returns_int() -> int:
+    return 42
+
+def returns_bool() -> bool:
+    return True
+
+x = returns_bool
+
+reveal_type(returns_bool() + returns_int())  # revealed: int
+# TODO should be float, need to check arg type and fall back
+reveal_type(returns_int() + 4.12)  # revealed: int
+reveal_type(4.2 + returns_bool())  # revealed: float
+```
+
+## With literal types
+
+When we have a literal type for one operand, we're able to fall back to the
+instance handling for its instance super-type.
+
+```py
+class A:
+    def __add__(self, other) -> A: return self
+    def __radd__(self, other) -> A: return self
+
+reveal_type(A() + 1)  # revealed: A
+reveal_type(1 + A())  # revealed: int
+reveal_type(A() + "foo")  # revealed: A
+# TODO should be A since `str.__add__` doesn't support A instances
+reveal_type("foo" + A())  # revealed: @Todo
+reveal_type(A() + b"foo")  # revealed: A
+reveal_type(b"foo" + A())  # revealed: bytes
+reveal_type(A() + ())  # revealed: A
+# TODO this should be A, since tuple's `__add__` doesn't support A instances
+reveal_type(() + A())  # revealed: @Todo
+```
+
+## Unsupported
+
+### Dunder on instance
+
+The magic method must be on the class, not just on the instance:
+
+```py
+def add_impl(self, other) -> int: return 1
+
+class A:
+    def __init__(self):
+        self.__add__ = add_impl
+
+a = A()
+
+# error: [operator]
+# revealed: Unknown
+reveal_type(a + 1)
+```
+
+### Missing dunder
+
+```py
+class A: pass
+
+# error: [operator]
+# revealed: Unknown
+reveal_type(A() + A())
+```
+
+### Wrong position
+
+A left-hand dunder method doesn't apply for the right-hand operand, or vice versa:
+
+```py
+class A:
+    def __add__(self, other) -> int: ...
+
+class B:
+    def __radd__(self, other) -> int: ...
+
+# error: [operator]
+# revealed: Unknown
+reveal_type(1 + A())
+# error: [operator]
+# revealed: Unknown
+reveal_type(B() + 1)
+```
+
+### Wrong type
+
+TODO: check signature and error if `other` is the wrong type

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

@@ -42,6 +42,6 @@ e = 1.0 / 0
 reveal_type(a)  # revealed: float
 reveal_type(b)  # revealed: int
 reveal_type(c)  # revealed: int
-reveal_type(d)  # revealed: @Todo
-reveal_type(e)  # revealed: @Todo
+reveal_type(d)  # revealed: float
+reveal_type(e)  # revealed: float
 ```

crates/red_knot_python_semantic/src/types/infer.rs

@@ -2489,6 +2489,15 @@ impl<'db> TypeInferenceBuilder<'db> {
             self.check_division_by_zero(binary, left_ty);
         }
 
+        self.infer_binary_expression_type(left_ty, right_ty, *op)
+    }
+
+    fn infer_binary_expression_type(
+        &mut self,
+        left_ty: Type<'db>,
+        right_ty: Type<'db>,
+        op: ast::Operator,
+    ) -> Type<'db> {
         match (left_ty, right_ty, op) {
             // When interacting with Todo, Any and Unknown should propagate (as if we fix this
             // `Todo` in the future, the result would then become Any or Unknown, respectively.)
@@ -2581,6 +2590,62 @@ impl<'db> TypeInferenceBuilder<'db> {
                 }
             }
 
+            (Type::Instance(_), Type::IntLiteral(_), op) => {
+                self.infer_binary_expression_type(left_ty, KnownClass::Int.to_instance(self.db), op)
+            }
+
+            (Type::IntLiteral(_), Type::Instance(_), op) => self.infer_binary_expression_type(
+                KnownClass::Int.to_instance(self.db),
+                right_ty,
+                op,
+            ),
+
+            (Type::Instance(_), Type::Tuple(_), op) => self.infer_binary_expression_type(
+                left_ty,
+                KnownClass::Tuple.to_instance(self.db),
+                op,
+            ),
+
+            (Type::Tuple(_), Type::Instance(_), op) => self.infer_binary_expression_type(
+                KnownClass::Tuple.to_instance(self.db),
+                right_ty,
+                op,
+            ),
+
+            (Type::Instance(_), Type::StringLiteral(_), op) => {
+                self.infer_binary_expression_type(left_ty, KnownClass::Str.to_instance(self.db), op)
+            }
+
+            (Type::StringLiteral(_), Type::Instance(_), op) => self.infer_binary_expression_type(
+                KnownClass::Str.to_instance(self.db),
+                right_ty,
+                op,
+            ),
+
+            (Type::Instance(_), Type::BytesLiteral(_), op) => self.infer_binary_expression_type(
+                left_ty,
+                KnownClass::Bytes.to_instance(self.db),
+                op,
+            ),
+
+            (Type::BytesLiteral(_), Type::Instance(_), op) => self.infer_binary_expression_type(
+                KnownClass::Bytes.to_instance(self.db),
+                right_ty,
+                op,
+            ),
+
+            (Type::Instance(left_class), Type::Instance(right_class), op) => left_class
+                .class_member(self.db, op.dunder())
+                .call(self.db, &[left_ty, right_ty])
+                .return_ty(self.db)
+                .unwrap_or_else(|| {
+                    right_class
+                        .class_member(self.db, op.reflected_dunder())
+                        .call(self.db, &[right_ty, left_ty])
+                        .return_ty(self.db)
+                        .unwrap_or(Type::Unknown)
+                }),
+
             _ => Type::Todo, // TODO
         }
     }

crates/ruff_python_ast/src/nodes.rs

@@ -2971,6 +2971,42 @@ impl Operator {
             Operator::FloorDiv => "//",
         }
     }
+
+    pub const fn dunder(self) -> &'static str {
+        match self {
+            Operator::Add => "__add__",
+            Operator::Sub => "__sub__",
+            Operator::Mult => "__mul__",
+            Operator::MatMult => "__matmul__",
+            Operator::Div => "__truediv__",
+            Operator::Mod => "__mod__",
+            Operator::Pow => "__pow__",
+            Operator::LShift => "__lshift__",
+            Operator::RShift => "__rshift__",
+            Operator::BitOr => "__or__",
+            Operator::BitXor => "__xor__",
+            Operator::BitAnd => "__and__",
+            Operator::FloorDiv => "__floordiv__",
+        }
+    }
+
+    pub const fn reflected_dunder(self) -> &'static str {
+        match self {
+            Operator::Add => "__radd__",
+            Operator::Sub => "__rsub__",
+            Operator::Mult => "__rmul__",
+            Operator::MatMult => "__rmatmul__",
+            Operator::Div => "__rtruediv__",
+            Operator::Mod => "__rmod__",
+            Operator::Pow => "__rpow__",
+            Operator::LShift => "__rlshift__",
+            Operator::RShift => "__rrshift__",
+            Operator::BitOr => "__ror__",
+            Operator::BitXor => "__rxor__",
+            Operator::BitAnd => "__rand__",
+            Operator::FloorDiv => "__rfloordiv__",
+        }
+    }
 }
 
 impl fmt::Display for Operator {