[ty] Infer definite equality comparison results

crates/ty_python_semantic/resources/mdtest/comparison/enums.md

@@ -19,3 +19,45 @@ reveal_type(Answer.NO is Answer.YES)  # revealed: Literal[False]
 reveal_type(Answer.NO is not Answer.NO)  # revealed: Literal[False]
 reveal_type(Answer.NO is not Answer.YES)  # revealed: Literal[True]
 ```
+
+Equality result inference uses the same runtime semantics as equality narrowing. This allows us to
+recognize when two enum domains cannot contain equal values, even if neither operand is a singleton:
+
+```py
+from enum import Enum
+from typing import Literal
+
+class Choice(str, Enum):
+    FIRST = "first"
+    SECOND = "second"
+    THIRD = "third"
+    FOURTH = "fourth"
+
+def compare(
+    left: Literal[Choice.FIRST, Choice.SECOND],
+    right: Literal[Choice.THIRD, Choice.FOURTH],
+):
+    reveal_type(left == right)  # revealed: Literal[False]
+    reveal_type(left != right)  # revealed: Literal[True]
+```
+
+When equality semantics are custom, result inference falls back to the corresponding dunder method:
+
+```py
+from enum import Enum
+
+class EqResult: ...
+class NeResult: ...
+
+class CustomEquality(Enum):
+    MEMBER = 1
+
+    def __eq__(self, other: object) -> EqResult:  # error: [invalid-method-override]
+        return EqResult()
+
+    def __ne__(self, other: object) -> NeResult:  # error: [invalid-method-override]
+        return NeResult()
+
+reveal_type(CustomEquality.MEMBER == CustomEquality.MEMBER)  # revealed: EqResult
+reveal_type(CustomEquality.MEMBER != CustomEquality.MEMBER)  # revealed: NeResult
+```

crates/ty_python_semantic/resources/mdtest/comparison/tuples.md

@@ -88,11 +88,9 @@ If two tuples contain types that do not support comparison, the result may be `U
 a = (1, 2)
 b = (1, "hello")
 
-# TODO: should be Literal[False], once we implement (in)equality for mismatched literals
-reveal_type(a == b)  # revealed: bool
+reveal_type(a == b)  # revealed: Literal[False]
 
-# TODO: should be Literal[True], once we implement (in)equality for mismatched literals
-reveal_type(a != b)  # revealed: bool
+reveal_type(a != b)  # revealed: Literal[True]
 
 # error: [unsupported-operator] "Operator `<` is not supported between objects of type `tuple[Literal[1], Literal[2]]` and `tuple[Literal[1], Literal["hello"]]`"
 reveal_type(a < b)  # revealed: Unknown
@@ -325,10 +323,8 @@ c = (1, 2, 3)
 reveal_type(a is (1, 2))  # revealed: bool
 reveal_type(a is not (1, 2))  # revealed: bool
 
-# TODO should be Literal[False] once we implement comparison of mismatched literal types
-reveal_type(a is b)  # revealed: bool
-# TODO should be Literal[True] once we implement comparison of mismatched literal types
-reveal_type(a is not b)  # revealed: bool
+reveal_type(a is b)  # revealed: Literal[False]
+reveal_type(a is not b)  # revealed: Literal[True]
 
 reveal_type(a is c)  # revealed: Literal[False]
 reveal_type(a is not c)  # revealed: Literal[True]

crates/ty_python_semantic/resources/mdtest/snapshots/tuples.md_-_Comparison___Tuples_-_Heterogeneous_-_Value_Comparisons_-_Comparison_Unsupport…_(966dd82bd3668d0e).snap

@@ -16,55 +16,53 @@ mdtest path: crates/ty_python_semantic/resources/mdtest/comparison/tuples.md
  1 | a = (1, 2)
  2 | b = (1, "hello")
  3 | 
- 4 | # TODO: should be Literal[False], once we implement (in)equality for mismatched literals
- 5 | reveal_type(a == b)  # revealed: bool
- 6 | 
- 7 | # TODO: should be Literal[True], once we implement (in)equality for mismatched literals
- 8 | reveal_type(a != b)  # revealed: bool
- 9 | 
-10 | # error: [unsupported-operator] "Operator `<` is not supported between objects of type `tuple[Literal[1], Literal[2]]` and `tuple[Literal[1], Literal["hello"]]`"
-11 | reveal_type(a < b)  # revealed: Unknown
-12 | # error: [unsupported-operator] "Operator `<=` is not supported between objects of type `tuple[Literal[1], Literal[2]]` and `tuple[Literal[1], Literal["hello"]]`"
-13 | reveal_type(a <= b)  # revealed: Unknown
-14 | # error: [unsupported-operator] "Operator `>` is not supported between objects of type `tuple[Literal[1], Literal[2]]` and `tuple[Literal[1], Literal["hello"]]`"
-15 | reveal_type(a > b)  # revealed: Unknown
-16 | # error: [unsupported-operator] "Operator `>=` is not supported between objects of type `tuple[Literal[1], Literal[2]]` and `tuple[Literal[1], Literal["hello"]]`"
-17 | reveal_type(a >= b)  # revealed: Unknown
-18 | # error: [unsupported-operator]
-19 | # error: [unsupported-operator]
-20 | reveal_type((object(),) < (object(),) < (object(),))  # revealed: Unknown
-21 | a = (1, 2)
-22 | b = (999999, "hello")
-23 | 
-24 | reveal_type(a == b)  # revealed: Literal[False]
-25 | reveal_type(a != b)  # revealed: Literal[True]
-26 | reveal_type(a < b)  # revealed: Literal[True]
-27 | reveal_type(a <= b)  # revealed: Literal[True]
-28 | reveal_type(a > b)  # revealed: Literal[False]
-29 | reveal_type(a >= b)  # revealed: Literal[False]
+ 4 | reveal_type(a == b)  # revealed: Literal[False]
+ 5 | 
+ 6 | reveal_type(a != b)  # revealed: Literal[True]
+ 7 | 
+ 8 | # error: [unsupported-operator] "Operator `<` is not supported between objects of type `tuple[Literal[1], Literal[2]]` and `tuple[Literal[1], Literal["hello"]]`"
+ 9 | reveal_type(a < b)  # revealed: Unknown
+10 | # error: [unsupported-operator] "Operator `<=` is not supported between objects of type `tuple[Literal[1], Literal[2]]` and `tuple[Literal[1], Literal["hello"]]`"
+11 | reveal_type(a <= b)  # revealed: Unknown
+12 | # error: [unsupported-operator] "Operator `>` is not supported between objects of type `tuple[Literal[1], Literal[2]]` and `tuple[Literal[1], Literal["hello"]]`"
+13 | reveal_type(a > b)  # revealed: Unknown
+14 | # error: [unsupported-operator] "Operator `>=` is not supported between objects of type `tuple[Literal[1], Literal[2]]` and `tuple[Literal[1], Literal["hello"]]`"
+15 | reveal_type(a >= b)  # revealed: Unknown
+16 | # error: [unsupported-operator]
+17 | # error: [unsupported-operator]
+18 | reveal_type((object(),) < (object(),) < (object(),))  # revealed: Unknown
+19 | a = (1, 2)
+20 | b = (999999, "hello")
+21 | 
+22 | reveal_type(a == b)  # revealed: Literal[False]
+23 | reveal_type(a != b)  # revealed: Literal[True]
+24 | reveal_type(a < b)  # revealed: Literal[True]
+25 | reveal_type(a <= b)  # revealed: Literal[True]
+26 | reveal_type(a > b)  # revealed: Literal[False]
+27 | reveal_type(a >= b)  # revealed: Literal[False]
 ```
 
 # Diagnostics
 
 ```
 error[unsupported-operator]: Unsupported `<` operation
-  --> src/mdtest_snippet.py:11:13
-   |
-11 | reveal_type(a < b)  # revealed: Unknown
-   |             -^^^-
-   |             |   |
-   |             |   Has type `tuple[Literal[1], Literal["hello"]]`
-   |             Has type `tuple[Literal[1], Literal[2]]`
-   |
+ --> src/mdtest_snippet.py:9:13
+  |
+9 | reveal_type(a < b)  # revealed: Unknown
+  |             -^^^-
+  |             |   |
+  |             |   Has type `tuple[Literal[1], Literal["hello"]]`
+  |             Has type `tuple[Literal[1], Literal[2]]`
+  |
 info: Operation fails because operator `<` is not supported between the tuple elements at index 2 (of type `Literal[2]` and `Literal["hello"]`)
 
 ```
 
 ```
 error[unsupported-operator]: Unsupported `<=` operation
-  --> src/mdtest_snippet.py:13:13
+  --> src/mdtest_snippet.py:11:13
    |
-13 | reveal_type(a <= b)  # revealed: Unknown
+11 | reveal_type(a <= b)  # revealed: Unknown
    |             -^^^^-
    |             |    |
    |             |    Has type `tuple[Literal[1], Literal["hello"]]`
@@ -76,9 +74,9 @@ info: Operation fails because operator `<=` is not supported between the tuple e
 
 ```
 error[unsupported-operator]: Unsupported `>` operation
-  --> src/mdtest_snippet.py:15:13
+  --> src/mdtest_snippet.py:13:13
    |
-15 | reveal_type(a > b)  # revealed: Unknown
+13 | reveal_type(a > b)  # revealed: Unknown
    |             -^^^-
    |             |   |
    |             |   Has type `tuple[Literal[1], Literal["hello"]]`
@@ -90,9 +88,9 @@ info: Operation fails because operator `>` is not supported between the tuple el
 
 ```
 error[unsupported-operator]: Unsupported `>=` operation
-  --> src/mdtest_snippet.py:17:13
+  --> src/mdtest_snippet.py:15:13
    |
-17 | reveal_type(a >= b)  # revealed: Unknown
+15 | reveal_type(a >= b)  # revealed: Unknown
    |             -^^^^-
    |             |    |
    |             |    Has type `tuple[Literal[1], Literal["hello"]]`
@@ -104,9 +102,9 @@ info: Operation fails because operator `>=` is not supported between the tuple e
 
 ```
 error[unsupported-operator]: Unsupported `<` operation
-  --> src/mdtest_snippet.py:20:13
+  --> src/mdtest_snippet.py:18:13
    |
-20 | reveal_type((object(),) < (object(),) < (object(),))  # revealed: Unknown
+18 | reveal_type((object(),) < (object(),) < (object(),))  # revealed: Unknown
    |             -----------^^^-----------
    |             |
    |             Both operands have type `tuple[object]`
@@ -117,9 +115,9 @@ info: Operation fails because operator `<` is not supported between the tuple el
 
 ```
 error[unsupported-operator]: Unsupported `<` operation
-  --> src/mdtest_snippet.py:20:27
+  --> src/mdtest_snippet.py:18:27
    |
-20 | reveal_type((object(),) < (object(),) < (object(),))  # revealed: Unknown
+18 | reveal_type((object(),) < (object(),) < (object(),))  # revealed: Unknown
    |                           -----------^^^-----------
    |                           |
    |                           Both operands have type `tuple[object]`

crates/ty_python_semantic/src/types/equality.rs

@@ -217,12 +217,27 @@ pub(crate) fn equality_truthiness<'db>(
     left: Type<'db>,
     right: Type<'db>,
 ) -> Truthiness {
-    match ComparisonEvaluator::new(db).evaluate(
-        left,
-        right,
-        ComparisonBranch::Positive,
-        ComparisonOperator::Equality,
-    ) {
+    comparison_truthiness(db, left, right, ComparisonOperator::Equality)
+}
+
+/// Return the truthiness of `left != right` when it is known for every represented runtime value.
+///
+/// A result that only permits narrowing remains ambiguous because it can still evaluate either way.
+pub(super) fn inequality_truthiness<'db>(
+    db: &'db dyn Db,
+    left: Type<'db>,
+    right: Type<'db>,
+) -> Truthiness {
+    comparison_truthiness(db, left, right, ComparisonOperator::Inequality)
+}
+
+fn comparison_truthiness<'db>(
+    db: &'db dyn Db,
+    left: Type<'db>,
+    right: Type<'db>,
+    operator: ComparisonOperator,
+) -> Truthiness {
+    match ComparisonEvaluator::new(db).evaluate(left, right, ComparisonBranch::Positive, operator) {
         ComparisonResult::AlwaysTrue => Truthiness::AlwaysTrue,
         ComparisonResult::AlwaysFalse => Truthiness::AlwaysFalse,
         ComparisonResult::CanNarrow(_) | ComparisonResult::Ambiguous => Truthiness::Ambiguous,

crates/ty_python_semantic/src/types/infer/comparisons.rs

@@ -7,6 +7,7 @@ use crate::types::call::{CallArguments, CallDunderError};
 use crate::types::constraints::ConstraintSetBuilder;
 use crate::types::context::InferContext;
 use crate::types::cyclic::CycleDetector;
+use crate::types::equality::{equality_truthiness, inequality_truthiness};
 use crate::types::tuple::TupleSpec;
 use crate::types::{
     DynamicType, IntersectionBuilder, IntersectionType, KnownClass, KnownInstanceType,
@@ -170,6 +171,15 @@ pub(super) fn infer_binary_type_comparison<'db>(
         }
     };
 
+    let comparison_truthiness = match op {
+        ast::CmpOp::Eq => equality_truthiness(db, left, right),
+        ast::CmpOp::NotEq => inequality_truthiness(db, left, right),
+        _ => Truthiness::Ambiguous,
+    };
+    if comparison_truthiness != Truthiness::Ambiguous {
+        return Ok(Type::from_truthiness(db, comparison_truthiness));
+    }
+
     let comparison_result = match (left, right) {
         (Type::EnumComplement(complement), right) => Some(infer_binary_type_comparison(
             context,
@@ -545,27 +555,6 @@ pub(super) fn infer_binary_type_comparison<'db>(
                     Some(Ok(result))
                 }
 
-                (LiteralValueTypeKind::Enum(literal_1), LiteralValueTypeKind::Enum(literal_2))
-                    if op == ast::CmpOp::Eq =>
-                {
-                    Some(Ok(
-                        match try_dunder(MemberLookupPolicy::MRO_NO_OBJECT_FALLBACK) {
-                            Ok(ty) => ty,
-                            Err(_) => Type::bool_literal(literal_1 == literal_2),
-                        },
-                    ))
-                }
-                (LiteralValueTypeKind::Enum(literal_1), LiteralValueTypeKind::Enum(literal_2))
-                    if op == ast::CmpOp::NotEq =>
-                {
-                    Some(Ok(
-                        match try_dunder(MemberLookupPolicy::MRO_NO_OBJECT_FALLBACK) {
-                            Ok(ty) => ty,
-                            Err(_) => Type::bool_literal(literal_1 != literal_2),
-                        },
-                    ))
-                }
-
                 _ => None,
             }
         }