Skip to content

Commit 270318c

Browse files
[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.)
1 parent d03b12e commit 270318c

File tree

3 files changed

+54
-16
lines changed

3 files changed

+54
-16
lines changed
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,22 @@
1+
# Binary operations on tuples
2+
3+
## Concatenation for heterogeneous tuples
4+
5+
```py
6+
reveal_type((1, 2) + (3, 4)) # revealed: tuple[Literal[1], Literal[2], Literal[3], Literal[4]]
7+
reveal_type(() + (1, 2)) # revealed: tuple[Literal[1], Literal[2]]
8+
reveal_type((1, 2) + ()) # revealed: tuple[Literal[1], Literal[2]]
9+
reveal_type(() + ()) # revealed: tuple[()]
10+
11+
def _(x: tuple[int, str], y: tuple[None, tuple[int]]):
12+
reveal_type(x + y) # revealed: tuple[int, str, None, tuple[int]]
13+
reveal_type(y + x) # revealed: tuple[None, tuple[int], int, str]
14+
```
15+
16+
## Concatenation for homogeneous tuples
17+
18+
```py
19+
def _(x: tuple[int, ...], y: tuple[str, ...]):
20+
reveal_type(x + y) # revealed: @Todo(full tuple[...] support)
21+
reveal_type(x + (1, 2)) # revealed: @Todo(full tuple[...] support)
22+
```

crates/red_knot_python_semantic/resources/mdtest/slots.md

+18-16
Original file line numberDiff line numberDiff line change
@@ -113,6 +113,24 @@ class D(B, A): ... # fine
113113
class E(B, C, A): ... # fine
114114
```
115115

116+
## Post-hoc modifications
117+
118+
```py
119+
class A:
120+
__slots__ = ()
121+
__slots__ += ("a", "b")
122+
123+
reveal_type(A.__slots__) # revealed: tuple[Literal["a"], Literal["b"]]
124+
125+
class B:
126+
__slots__ = ("c", "d")
127+
128+
class C(
129+
A, # error: [incompatible-slots]
130+
B, # error: [incompatible-slots]
131+
): ...
132+
```
133+
116134
## False negatives
117135

118136
### Possibly unbound
@@ -160,22 +178,6 @@ class B:
160178
class C(A, B): ...
161179
```
162180

163-
### Post-hoc modifications
164-
165-
```py
166-
class A:
167-
__slots__ = ()
168-
__slots__ += ("a", "b")
169-
170-
reveal_type(A.__slots__) # revealed: @Todo(return type of decorated function)
171-
172-
class B:
173-
__slots__ = ("c", "d")
174-
175-
# False negative: [incompatible-slots]
176-
class C(A, B): ...
177-
```
178-
179181
### Built-ins with implicit layouts
180182

181183
```py

crates/red_knot_python_semantic/src/types/infer.rs

+14
Original file line numberDiff line numberDiff line change
@@ -4633,6 +4633,20 @@ impl<'db> TypeInferenceBuilder<'db> {
46334633
self.infer_binary_expression_type(left, Type::IntLiteral(i64::from(bool_value)), op)
46344634
}
46354635

4636+
(Type::Tuple(lhs), Type::Tuple(rhs), ast::Operator::Add) => {
4637+
// Note: this only works on heterogeneous tuples.
4638+
let lhs_elements = lhs.elements(self.db());
4639+
let rhs_elements = rhs.elements(self.db());
4640+
4641+
Some(TupleType::from_elements(
4642+
self.db(),
4643+
lhs_elements
4644+
.iter()
4645+
.copied()
4646+
.chain(rhs_elements.iter().copied()),
4647+
))
4648+
}
4649+
46364650
// We've handled all of the special cases that we support for literals, so we need to
46374651
// fall back on looking for dunder methods on one of the operand types.
46384652
(

0 commit comments

Comments
 (0)