From 270318c2e023aaf4f04aa56ea2bf6f8d426b16b7 Mon Sep 17 00:00:00 2001
From: cake-monotone <cake.monotone@gmail.com>
Date: Fri, 14 Mar 2025 20:29:57 +0900
Subject: [PATCH] [red-knot] fix: improve type inference for binary ops on
 tuples (#16725)

## Summary

This PR includes minor improvements to binary operation inference,
specifically for tuple concatenation.

### Before

```py
reveal_type((1, 2) + (3, 4))  # revealed: @Todo(return type of decorated function)
# If TODO is ignored, the revealed type would be `tuple[1|2|3|4, ...]`
```

The `builtins.tuple` type stub defines `__add__`, but it appears to only
work for homogeneous tuples. However, I think this limitation is not
ideal for many use cases.

### After

```py
reveal_type((1, 2) + (3, 4))  # revealed: tuple[Literal[1], Literal[2], Literal[3], Literal[4]]
```

## Test Plan

### Added
- `mdtest/binary/tuples.md`

### Affected
- `mdtest/slots.md` (a test have been moved out of the `False-Negative`
block.)
---
 .../resources/mdtest/binary/tuples.md         | 22 ++++++++++++
 .../resources/mdtest/slots.md                 | 34 ++++++++++---------
 .../src/types/infer.rs                        | 14 ++++++++
 3 files changed, 54 insertions(+), 16 deletions(-)
 create mode 100644 crates/red_knot_python_semantic/resources/mdtest/binary/tuples.md

diff --git a/crates/red_knot_python_semantic/resources/mdtest/binary/tuples.md b/crates/red_knot_python_semantic/resources/mdtest/binary/tuples.md
new file mode 100644
index 0000000000..49fb5b7333
--- /dev/null
+++ b/crates/red_knot_python_semantic/resources/mdtest/binary/tuples.md
@@ -0,0 +1,22 @@
+# Binary operations on tuples
+
+## Concatenation for heterogeneous tuples
+
+```py
+reveal_type((1, 2) + (3, 4))  # revealed: tuple[Literal[1], Literal[2], Literal[3], Literal[4]]
+reveal_type(() + (1, 2))  # revealed: tuple[Literal[1], Literal[2]]
+reveal_type((1, 2) + ())  # revealed: tuple[Literal[1], Literal[2]]
+reveal_type(() + ())  # revealed: tuple[()]
+
+def _(x: tuple[int, str], y: tuple[None, tuple[int]]):
+    reveal_type(x + y)  # revealed: tuple[int, str, None, tuple[int]]
+    reveal_type(y + x)  # revealed: tuple[None, tuple[int], int, str]
+```
+
+## Concatenation for homogeneous tuples
+
+```py
+def _(x: tuple[int, ...], y: tuple[str, ...]):
+    reveal_type(x + y)  # revealed: @Todo(full tuple[...] support)
+    reveal_type(x + (1, 2))  # revealed: @Todo(full tuple[...] support)
+```
diff --git a/crates/red_knot_python_semantic/resources/mdtest/slots.md b/crates/red_knot_python_semantic/resources/mdtest/slots.md
index f319d163c2..05e44eb4a3 100644
--- a/crates/red_knot_python_semantic/resources/mdtest/slots.md
+++ b/crates/red_knot_python_semantic/resources/mdtest/slots.md
@@ -113,6 +113,24 @@ class D(B, A): ...  # fine
 class E(B, C, A): ...  # fine
 ```
 
+## Post-hoc modifications
+
+```py
+class A:
+    __slots__ = ()
+    __slots__ += ("a", "b")
+
+reveal_type(A.__slots__)  # revealed: tuple[Literal["a"], Literal["b"]]
+
+class B:
+    __slots__ = ("c", "d")
+
+class C(
+    A,  # error: [incompatible-slots]
+    B,  # error: [incompatible-slots]
+): ...
+```
+
 ## False negatives
 
 ### Possibly unbound
@@ -160,22 +178,6 @@ class B:
 class C(A, B): ...
 ```
 
-### Post-hoc modifications
-
-```py
-class A:
-    __slots__ = ()
-    __slots__ += ("a", "b")
-
-reveal_type(A.__slots__)  # revealed: @Todo(return type of decorated function)
-
-class B:
-    __slots__ = ("c", "d")
-
-# False negative: [incompatible-slots]
-class C(A, B): ...
-```
-
 ### Built-ins with implicit layouts
 
 ```py
diff --git a/crates/red_knot_python_semantic/src/types/infer.rs b/crates/red_knot_python_semantic/src/types/infer.rs
index d216ccd840..cec1cb2180 100644
--- a/crates/red_knot_python_semantic/src/types/infer.rs
+++ b/crates/red_knot_python_semantic/src/types/infer.rs
@@ -4633,6 +4633,20 @@ impl<'db> TypeInferenceBuilder<'db> {
                 self.infer_binary_expression_type(left, Type::IntLiteral(i64::from(bool_value)), op)
             }
 
+            (Type::Tuple(lhs), Type::Tuple(rhs), ast::Operator::Add) => {
+                // Note: this only works on heterogeneous tuples.
+                let lhs_elements = lhs.elements(self.db());
+                let rhs_elements = rhs.elements(self.db());
+
+                Some(TupleType::from_elements(
+                    self.db(),
+                    lhs_elements
+                        .iter()
+                        .copied()
+                        .chain(rhs_elements.iter().copied()),
+                ))
+            }
+
             // We've handled all of the special cases that we support for literals, so we need to
             // fall back on looking for dunder methods on one of the operand types.
             (