[red-knot] Inference for comparison of union types (#13781)

## Summary

Add type inference for comparisons involving union types. For example:
```py
one_or_two = 1 if flag else 2

reveal_type(one_or_two <= 2)  # revealed: Literal[True]
reveal_type(one_or_two <= 1)  # revealed: bool
reveal_type(one_or_two <= 0)  # revealed: Literal[False]
```

closes #13779

## Test Plan

See `resources/mdtest/comparison/unions.md`
This commit is contained in:
David Peter 2024-10-17 11:03:37 +02:00 committed by GitHub
parent 6b7a738825
commit 5c3c0c4705
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
4 changed files with 108 additions and 2 deletions

View file

@ -0,0 +1,76 @@
# Comparison: Unions
## Union on one side of the comparison
Comparisons on union types need to consider all possible cases:
```py
one_or_two = 1 if flag else 2
reveal_type(one_or_two <= 2) # revealed: Literal[True]
reveal_type(one_or_two <= 1) # revealed: bool
reveal_type(one_or_two <= 0) # revealed: Literal[False]
reveal_type(2 >= one_or_two) # revealed: Literal[True]
reveal_type(1 >= one_or_two) # revealed: bool
reveal_type(0 >= one_or_two) # revealed: Literal[False]
reveal_type(one_or_two < 1) # revealed: Literal[False]
reveal_type(one_or_two < 2) # revealed: bool
reveal_type(one_or_two < 3) # revealed: Literal[True]
reveal_type(one_or_two > 0) # revealed: Literal[True]
reveal_type(one_or_two > 1) # revealed: bool
reveal_type(one_or_two > 2) # revealed: Literal[False]
reveal_type(one_or_two == 3) # revealed: Literal[False]
reveal_type(one_or_two == 1) # revealed: bool
reveal_type(one_or_two != 3) # revealed: Literal[True]
reveal_type(one_or_two != 1) # revealed: bool
a_or_ab = "a" if flag else "ab"
reveal_type(a_or_ab in "ab") # revealed: Literal[True]
reveal_type("a" in a_or_ab) # revealed: Literal[True]
reveal_type("c" not in a_or_ab) # revealed: Literal[True]
reveal_type("a" not in a_or_ab) # revealed: Literal[False]
reveal_type("b" in a_or_ab) # revealed: bool
reveal_type("b" not in a_or_ab) # revealed: bool
one_or_none = 1 if flag else None
reveal_type(one_or_none is None) # revealed: bool
reveal_type(one_or_none is not None) # revealed: bool
```
## Union on both sides of the comparison
With unions on both sides, we need to consider the full cross product of
options when building the resulting (union) type:
```py
small = 1 if flag_s else 2
large = 2 if flag_l else 3
reveal_type(small <= large) # revealed: Literal[True]
reveal_type(small >= large) # revealed: bool
reveal_type(small < large) # revealed: bool
reveal_type(small > large) # revealed: Literal[False]
```
## Unsupported operations
Make sure we emit a diagnostic if *any* of the possible comparisons is
unsupported. For now, we fall back to `bool` for the result type instead of
trying to infer something more precise from the other (supported) variants:
```py
x = [1, 2] if flag else 1
result = 1 in x # error: "Operator `in` is not supported"
reveal_type(result) # revealed: bool
```