From 860b95a318d5dabb9b33e5595b8cc229d0841ce2 Mon Sep 17 00:00:00 2001
From: David Peter <sharkdp@users.noreply.github.com>
Date: Wed, 12 Mar 2025 08:21:54 +0100
Subject: [PATCH] [red-knot] Binary operator inference for union types (#16601)

## Summary

Properly handle binary operator inference for union types.

This fixes a bug I noticed while looking at ecosystem results. The MRE
version of it is this:

```py
def sub(x: float, y: float):
    # Red Knot: Operator `-` is unsupported between objects of type `int | float` and `int | float`
    return x - y
```

## Test Plan

- New Markdown tests.
- Expected diff in the ecosystem checks
---
 .../resources/mdtest/binary/unions.md         | 51 +++++++++++++++++++
 .../src/types/infer.rs                        | 19 ++++++-
 2 files changed, 68 insertions(+), 2 deletions(-)
 create mode 100644 crates/red_knot_python_semantic/resources/mdtest/binary/unions.md

diff --git a/crates/red_knot_python_semantic/resources/mdtest/binary/unions.md b/crates/red_knot_python_semantic/resources/mdtest/binary/unions.md
new file mode 100644
index 0000000000..0f5fb09bf5
--- /dev/null
+++ b/crates/red_knot_python_semantic/resources/mdtest/binary/unions.md
@@ -0,0 +1,51 @@
+# Binary operations on union types
+
+Binary operations on union types are only available if they are supported for all possible
+combinations of types:
+
+```py
+def f1(i: int, u: int | None):
+    # error: [unsupported-operator] "Operator `+` is unsupported between objects of type `int` and `int | None`"
+    reveal_type(i + u)  # revealed: Unknown
+    # error: [unsupported-operator] "Operator `+` is unsupported between objects of type `int | None` and `int`"
+    reveal_type(u + i)  # revealed: Unknown
+```
+
+`int` can be added to `int`, and `str` can be added to `str`, but expressions of type `int | str`
+cannot be added, because that would require addition of `int` and `str` or vice versa:
+
+```py
+def f2(i: int, s: str, int_or_str: int | str):
+    i + i
+    s + s
+    # error: [unsupported-operator] "Operator `+` is unsupported between objects of type `int | str` and `int | str`"
+    reveal_type(int_or_str + int_or_str)  # revealed: Unknown
+```
+
+However, if an operation is supported for all possible combinations, the result will be a union of
+the possible outcomes:
+
+```py
+from typing import Literal
+
+def f3(two_or_three: Literal[2, 3], a_or_b: Literal["a", "b"]):
+    reveal_type(two_or_three + two_or_three)  # revealed: Literal[4, 5, 6]
+    reveal_type(two_or_three**two_or_three)  # revealed: Literal[4, 8, 9, 27]
+
+    reveal_type(a_or_b + a_or_b)  # revealed: Literal["aa", "ab", "ba", "bb"]
+
+    reveal_type(two_or_three * a_or_b)  # revealed: Literal["aa", "bb", "aaa", "bbb"]
+```
+
+We treat a type annotation of `float` as a union of `int` and `float`, so union handling is relevant
+here:
+
+```py
+def f4(x: float, y: float):
+    reveal_type(x + y)  # revealed: int | float
+    reveal_type(x - y)  # revealed: int | float
+    reveal_type(x * y)  # revealed: int | float
+    reveal_type(x / y)  # revealed: int | float
+    reveal_type(x // y)  # revealed: int | float
+    reveal_type(x % y)  # revealed: int | float
+```
diff --git a/crates/red_knot_python_semantic/src/types/infer.rs b/crates/red_knot_python_semantic/src/types/infer.rs
index bb28029c21..8a64d0e053 100644
--- a/crates/red_knot_python_semantic/src/types/infer.rs
+++ b/crates/red_knot_python_semantic/src/types/infer.rs
@@ -4144,6 +4144,23 @@ impl<'db> TypeInferenceBuilder<'db> {
         op: ast::Operator,
     ) -> Option<Type<'db>> {
         match (left_ty, right_ty, op) {
+            (Type::Union(lhs_union), rhs, _) => {
+                let mut union = UnionBuilder::new(self.db());
+                for lhs in lhs_union.elements(self.db()) {
+                    let result = self.infer_binary_expression_type(*lhs, rhs, op)?;
+                    union = union.add(result);
+                }
+                Some(union.build())
+            }
+            (lhs, Type::Union(rhs_union), _) => {
+                let mut union = UnionBuilder::new(self.db());
+                for rhs in rhs_union.elements(self.db()) {
+                    let result = self.infer_binary_expression_type(lhs, *rhs, op)?;
+                    union = union.add(result);
+                }
+                Some(union.build())
+            }
+
             // Non-todo Anys take precedence over Todos (as if we fix this `Todo` in the future,
             // the result would then become Any or Unknown, respectively).
             (any @ Type::Dynamic(DynamicType::Any), _, _)
@@ -4275,7 +4292,6 @@ impl<'db> TypeInferenceBuilder<'db> {
                 | Type::SubclassOf(_)
                 | Type::Instance(_)
                 | Type::KnownInstance(_)
-                | Type::Union(_)
                 | Type::Intersection(_)
                 | Type::AlwaysTruthy
                 | Type::AlwaysFalsy
@@ -4292,7 +4308,6 @@ impl<'db> TypeInferenceBuilder<'db> {
                 | Type::SubclassOf(_)
                 | Type::Instance(_)
                 | Type::KnownInstance(_)
-                | Type::Union(_)
                 | Type::Intersection(_)
                 | Type::AlwaysTruthy
                 | Type::AlwaysFalsy