[red-knot] Refactoring the inference logic of lexicographic comparisons (#14422)

## Summary

closes #14279

### Limitations of the Current Implementation
#### Incorrect Error Propagation

In the current implementation of lexicographic comparisons, if the
result of an Eq operation is Ambiguous, the comparison stops
immediately, returning a bool instance. While this may yield correct
inferences, it fails to capture unsupported-operation errors that might
occur in subsequent comparisons.
```py
class A: ...

(int_instance(), A()) < (int_instance(), A())  # should error
```

#### Weak Inference in Specific Cases

> Example: `(int_instance(), "foo") == (int_instance(), "bar")`
> Current result: `bool`
> Expected result: `Literal[False]`

`Eq` and `NotEq` have unique behavior in lexicographic comparisons
compared to other operators. Specifically:
- For `Eq`, if any non-equal pair exists within the tuples being
compared, we can immediately conclude that the tuples are not equal.
- For `NotEq`, if any equal pair exists, we can conclude that the tuples
are unequal.

```py
a = (str_instance(), int_instance(), "foo")

reveal_type(a == a)  # revealed: bool
reveal_type(a != a)  # revealed: bool

b = (str_instance(), int_instance(), "bar")

reveal_type(a == b)  # revealed: bool  # should be Literal[False]
reveal_type(a != b)  # revealed: bool  # should be Literal[True]
```
#### Incorrect Support for Non-Boolean Rich Comparisons

In CPython, aside from `==` and `!=`, tuple comparisons return a
non-boolean result as-is. Tuples do not convert the value into `bool`.

Note: If all pairwise `==` comparisons between elements in the tuples
return Truthy, the comparison then considers the tuples' lengths.
Regardless of the return type of the dunder methods, the final result
can still be a boolean.

```py
from __future__ import annotations

class A:
    def __eq__(self, o: object) -> str:
        return "hello"

    def __ne__(self, o: object) -> bytes:
        return b"world"

    def __lt__(self, o: A) -> float:
        return 3.14

a = (A(), A())

reveal_type(a == a)  # revealed: bool
reveal_type(a != a)  # revealed: bool
reveal_type(a < a)  # revealed: bool # should be: `float | Literal[False]`

```

### Key Changes
One of the major changes is that comparisons no longer end with a `bool`
result when a pairwise `Eq` result is `Ambiguous`. Instead, the function
attempts to infer all possible cases and unions the results. This
improvement allows for more robust type inference and better error
detection.

Additionally, as the function is now optimized for tuple comparisons,
the name has been changed from the more general
`infer_lexicographic_comparison` to `infer_tuple_rich_comparison`.

## Test Plan

mdtest included
This commit is contained in:
cake-monotone 2024-11-20 10:32:43 +09:00 committed by GitHub
parent 42c35b6f44
commit 6a4d207db7
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
4 changed files with 208 additions and 40 deletions

View file

@ -58,7 +58,9 @@ reveal_type(c >= d) # revealed: Literal[True]
#### Results with Ambiguity
```py
def bool_instance() -> bool: ...
def bool_instance() -> bool:
return True
def int_instance() -> int:
return 42
@ -134,23 +136,158 @@ reveal_type(c >= c) # revealed: Literal[True]
#### Non Boolean Rich Comparisons
Rich comparison methods defined in a class affect tuple comparisons as well. Proper type inference
should be possible even in cases where these methods return non-boolean types.
Note: Tuples use lexicographic comparisons. If the `==` result for all paired elements in the tuple
is True, the comparison then considers the tuples length. Regardless of the return type of the
dunder methods, the final result can still be a boolean value.
(+cpython: For tuples, `==` and `!=` always produce boolean results, regardless of the return type
of the dunder methods.)
```py
from __future__ import annotations
class A:
def __eq__(self, o) -> str: ...
def __ne__(self, o) -> int: ...
def __lt__(self, o) -> float: ...
def __le__(self, o) -> object: ...
def __gt__(self, o) -> tuple: ...
def __ge__(self, o) -> list: ...
def __eq__(self, o: object) -> str:
return "hello"
def __ne__(self, o: object) -> bytes:
return b"world"
def __lt__(self, o: A) -> float:
return 3.14
def __le__(self, o: A) -> complex:
return complex(0.5, -0.5)
def __gt__(self, o: A) -> tuple:
return (1, 2, 3)
def __ge__(self, o: A) -> list:
return [1, 2, 3]
a = (A(), A())
reveal_type(a == a) # revealed: bool
reveal_type(a != a) # revealed: bool
reveal_type(a < a) # revealed: float | Literal[False]
reveal_type(a <= a) # revealed: complex | Literal[True]
reveal_type(a > a) # revealed: tuple | Literal[False]
reveal_type(a >= a) # revealed: list | Literal[True]
# If lexicographic comparison is finished before comparing A()
b = ("1_foo", A())
c = ("2_bar", A())
reveal_type(b == c) # revealed: Literal[False]
reveal_type(b != c) # revealed: Literal[True]
reveal_type(b < c) # revealed: Literal[True]
reveal_type(b <= c) # revealed: Literal[True]
reveal_type(b > c) # revealed: Literal[False]
reveal_type(b >= c) # revealed: Literal[False]
class B:
def __lt__(self, o: B) -> set:
return set()
reveal_type((A(), B()) < (A(), B())) # revealed: float | set | Literal[False]
```
#### Special Handling of Eq and NotEq in Lexicographic Comparisons
> Example: `(int_instance(), "foo") == (int_instance(), "bar")`
`Eq` and `NotEq` have unique behavior compared to other operators in lexicographic comparisons.
Specifically, for `Eq`, if any non-equal pair exists within the tuples being compared, we can
immediately conclude that the tuples are not equal. Conversely, for `NotEq`, if any non-equal pair
exists, we can determine that the tuples are unequal.
In contrast, with operators like `<` and `>`, the comparison must consider each pair of elements
sequentially, and the final outcome might remain ambiguous until all pairs are compared.
```py
def str_instance() -> str:
return "hello"
def int_instance() -> int:
return 42
reveal_type("foo" == "bar") # revealed: Literal[False]
reveal_type(("foo",) == ("bar",)) # revealed: Literal[False]
reveal_type((4, "foo") == (4, "bar")) # revealed: Literal[False]
reveal_type((int_instance(), "foo") == (int_instance(), "bar")) # revealed: Literal[False]
a = (str_instance(), int_instance(), "foo")
reveal_type(a == a) # revealed: bool
reveal_type(a != a) # revealed: bool
reveal_type(a < a) # revealed: bool
reveal_type(a <= a) # revealed: bool
reveal_type(a > a) # revealed: bool
reveal_type(a >= a) # revealed: bool
b = (str_instance(), int_instance(), "bar")
reveal_type(a == b) # revealed: Literal[False]
reveal_type(a != b) # revealed: Literal[True]
reveal_type(a < b) # revealed: bool
reveal_type(a <= b) # revealed: bool
reveal_type(a > b) # revealed: bool
reveal_type(a >= b) # revealed: bool
c = (str_instance(), int_instance(), "foo", "different_length")
reveal_type(a == c) # revealed: Literal[False]
reveal_type(a != c) # revealed: Literal[True]
reveal_type(a < c) # revealed: bool
reveal_type(a <= c) # revealed: bool
reveal_type(a > c) # revealed: bool
reveal_type(a >= c) # revealed: bool
```
#### Error Propagation
Errors occurring within a tuple comparison should propagate outward. However, if the tuple
comparison can clearly conclude before encountering an error, the error should not be raised.
```py
def int_instance() -> int:
return 42
def str_instance() -> str:
return "hello"
class A: ...
# error: [unsupported-operator] "Operator `<` is not supported for types `A` and `A`"
A() < A()
# error: [unsupported-operator] "Operator `<=` is not supported for types `A` and `A`"
A() <= A()
# error: [unsupported-operator] "Operator `>` is not supported for types `A` and `A`"
A() > A()
# error: [unsupported-operator] "Operator `>=` is not supported for types `A` and `A`"
A() >= A()
a = (0, int_instance(), A())
# error: [unsupported-operator] "Operator `<` is not supported for types `A` and `A`, in comparing `tuple[Literal[0], int, A]` with `tuple[Literal[0], int, A]`"
reveal_type(a < a) # revealed: Unknown
# error: [unsupported-operator] "Operator `<=` is not supported for types `A` and `A`, in comparing `tuple[Literal[0], int, A]` with `tuple[Literal[0], int, A]`"
reveal_type(a <= a) # revealed: Unknown
# error: [unsupported-operator] "Operator `>` is not supported for types `A` and `A`, in comparing `tuple[Literal[0], int, A]` with `tuple[Literal[0], int, A]`"
reveal_type(a > a) # revealed: Unknown
# error: [unsupported-operator] "Operator `>=` is not supported for types `A` and `A`, in comparing `tuple[Literal[0], int, A]` with `tuple[Literal[0], int, A]`"
reveal_type(a >= a) # revealed: Unknown
# Comparison between `a` and `b` should only involve the first elements, `Literal[0]` and `Literal[99999]`,
# and should terminate immediately.
b = (99999, int_instance(), A())
reveal_type(a < b) # revealed: Literal[True]
reveal_type(a <= b) # revealed: Literal[True]
reveal_type(a > b) # revealed: Literal[False]
reveal_type(a >= b) # revealed: Literal[False]
```
### Membership Test Comparisons

View file

@ -4,6 +4,8 @@
def bool_instance() -> bool:
return True
class A: ...
a = 1 in 7 # error: "Operator `in` is not supported for types `Literal[1]` and `Literal[7]`"
reveal_type(a) # revealed: bool
@ -33,4 +35,8 @@ reveal_type(e) # revealed: bool
f = (1, 2) < (1, "hello")
# TODO: should be Unknown, once operand type check is implemented
reveal_type(f) # revealed: bool
# error: [unsupported-operator] "Operator `<` is not supported for types `A` and `A`, in comparing `tuple[bool, A]` with `tuple[bool, A]`"
g = (bool_instance(), A()) < (bool_instance(), A())
reveal_type(g) # revealed: Unknown
```

View file

@ -58,8 +58,7 @@ reveal_type(sys.version_info >= (3, 9, 1, "final", 0)) # revealed: bool
# emitting a lint diagnostic of some kind warning them about the probable error?
reveal_type(sys.version_info >= (3, 9, 1, "final", 0, 5)) # revealed: bool
# TODO: this should be `Literal[False]`; see #14279
reveal_type(sys.version_info == (3, 9, 1, "finallllll", 0)) # revealed: bool
reveal_type(sys.version_info == (3, 8, 1, "finallllll", 0)) # revealed: Literal[False]
```
## Imports and aliases

View file

@ -3643,16 +3643,16 @@ impl<'db> TypeInferenceBuilder<'db> {
let lhs_elements = lhs.elements(self.db);
let rhs_elements = rhs.elements(self.db);
let mut lexicographic_type_comparison =
|op| self.infer_lexicographic_type_comparison(lhs_elements, op, rhs_elements);
let mut tuple_rich_comparison =
|op| self.infer_tuple_rich_comparison(lhs_elements, op, rhs_elements);
match op {
ast::CmpOp::Eq => lexicographic_type_comparison(RichCompareOperator::Eq),
ast::CmpOp::NotEq => lexicographic_type_comparison(RichCompareOperator::Ne),
ast::CmpOp::Lt => lexicographic_type_comparison(RichCompareOperator::Lt),
ast::CmpOp::LtE => lexicographic_type_comparison(RichCompareOperator::Le),
ast::CmpOp::Gt => lexicographic_type_comparison(RichCompareOperator::Gt),
ast::CmpOp::GtE => lexicographic_type_comparison(RichCompareOperator::Ge),
ast::CmpOp::Eq => tuple_rich_comparison(RichCompareOperator::Eq),
ast::CmpOp::NotEq => tuple_rich_comparison(RichCompareOperator::Ne),
ast::CmpOp::Lt => tuple_rich_comparison(RichCompareOperator::Lt),
ast::CmpOp::LtE => tuple_rich_comparison(RichCompareOperator::Le),
ast::CmpOp::Gt => tuple_rich_comparison(RichCompareOperator::Gt),
ast::CmpOp::GtE => tuple_rich_comparison(RichCompareOperator::Ge),
ast::CmpOp::In | ast::CmpOp::NotIn => {
let mut eq_count = 0usize;
let mut not_eq_count = 0usize;
@ -3685,8 +3685,7 @@ impl<'db> TypeInferenceBuilder<'db> {
ast::CmpOp::Is | ast::CmpOp::IsNot => {
// - `[ast::CmpOp::Is]`: returns `false` if the elements are definitely unequal, otherwise `bool`
// - `[ast::CmpOp::IsNot]`: returns `true` if the elements are definitely unequal, otherwise `bool`
let eq_result = lexicographic_type_comparison(RichCompareOperator::Eq)
.expect(
let eq_result = tuple_rich_comparison(RichCompareOperator::Eq).expect(
"infer_binary_type_comparison should never return None for `CmpOp::Eq`",
);
@ -3751,53 +3750,80 @@ impl<'db> TypeInferenceBuilder<'db> {
}
}
/// Performs lexicographic comparison between two slices of types.
/// Simulates rich comparison between tuples and returns the inferred result.
/// This performs a lexicographic comparison, returning a union of all possible return types that could result from the comparison.
///
/// For lexicographic comparison, elements from both slices are compared pairwise using
/// `infer_binary_type_comparison`. If a conclusive result cannot be determined as a `BooleanLiteral`,
/// it returns `bool`. Returns `None` if the comparison is not supported.
fn infer_lexicographic_type_comparison(
/// basically it's based on cpython's `tuple_richcompare`
/// see `<https://github.com/python/cpython/blob/9d6366b60d01305fc5e45100e0cd13e358aa397d/Objects/tupleobject.c#L637>`
fn infer_tuple_rich_comparison(
&mut self,
left: &[Type<'db>],
op: RichCompareOperator,
right: &[Type<'db>],
) -> Result<Type<'db>, CompareUnsupportedError<'db>> {
// Compare paired elements from left and right slices
for (l_ty, r_ty) in left.iter().copied().zip(right.iter().copied()) {
let eq_result = self
let left_iter = left.iter().copied();
let right_iter = right.iter().copied();
let mut builder = UnionBuilder::new(self.db);
for (l_ty, r_ty) in left_iter.zip(right_iter) {
let pairwise_eq_result = self
.infer_binary_type_comparison(l_ty, ast::CmpOp::Eq, r_ty)
.expect("infer_binary_type_comparison should never return None for `CmpOp::Eq`");
match eq_result {
match pairwise_eq_result {
// If propagation is required, return the result as is
Type::Todo => return Ok(Type::Todo),
ty => match ty.bool(self.db) {
// Types are equal, continue to the next pair
// - AlwaysTrue : Continue to the next pair for lexicographic comparison
Truthiness::AlwaysTrue => continue,
// Types are not equal, perform the specified comparison and return the result
Truthiness::AlwaysFalse => {
return self.infer_binary_type_comparison(l_ty, op.into(), r_ty)
// - AlwaysFalse:
// Lexicographic comparisons will always terminate with this pair.
// Complete the comparison and return the result.
// - Ambiguous:
// Lexicographic comparisons might continue to the next pair (if eq_result is true),
// or terminate here (if eq_result is false).
// To account for cases where the comparison terminates here, add the pairwise comparison result to the union builder.
eq_truthiness @ (Truthiness::AlwaysFalse | Truthiness::Ambiguous) => {
let pairwise_compare_result = match op {
RichCompareOperator::Lt
| RichCompareOperator::Le
| RichCompareOperator::Gt
| RichCompareOperator::Ge => {
self.infer_binary_type_comparison(l_ty, op.into(), r_ty)?
}
// For `==` and `!=`, we already figure out the result from `pairwise_eq_result`
// NOTE: The CPython implementation does not account for non-boolean return types
// or cases where `!=` is not the negation of `==`, we also do not consider these cases.
RichCompareOperator::Eq => Type::BooleanLiteral(false),
RichCompareOperator::Ne => Type::BooleanLiteral(true),
};
builder = builder.add(pairwise_compare_result);
if eq_truthiness.is_ambiguous() {
continue;
}
return Ok(builder.build());
}
// If the intermediate result is ambiguous, we cannot determine the final result as BooleanLiteral.
// In this case, we simply return a bool instance.
Truthiness::Ambiguous => return Ok(KnownClass::Bool.to_instance(self.db)),
},
}
}
// At this point, the lengths of the two slices may be different, but the prefix of
// left and right slices is entirely identical.
// We return a comparison of the slice lengths based on the operator.
// if no more items to compare, we just compare sizes
let (left_len, right_len) = (left.len(), right.len());
Ok(Type::BooleanLiteral(match op {
builder = builder.add(Type::BooleanLiteral(match op {
RichCompareOperator::Eq => left_len == right_len,
RichCompareOperator::Ne => left_len != right_len,
RichCompareOperator::Lt => left_len < right_len,
RichCompareOperator::Le => left_len <= right_len,
RichCompareOperator::Gt => left_len > right_len,
RichCompareOperator::Ge => left_len >= right_len,
}))
}));
Ok(builder.build())
}
fn infer_subscript_expression(&mut self, subscript: &ast::ExprSubscript) -> Type<'db> {