[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
This commit is contained in:
David Peter 2025-03-12 08:21:54 +01:00 committed by GitHub
parent 6de2b2873b
commit 860b95a318
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
2 changed files with 68 additions and 2 deletions

View file

@ -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
```

View file

@ -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