Fix equality narrowing

Author: sharkdp

crates/ty_python_semantic/resources/mdtest/narrow/conditionals/eq.md

@@ -37,8 +37,7 @@ def _(answer: Answer):
     if answer != Answer.NO:
         reveal_type(answer)  # revealed: Literal[Answer.YES]
     else:
-        # TODO: This should be `Literal[Answer.NO]`
-        reveal_type(answer)  # revealed: Answer
+        reveal_type(answer)  # revealed: Literal[Answer.NO]
 ```
 
 This narrowing behavior is only safe if the enum has no custom `__eq__`/`__ne__` method:

crates/ty_python_semantic/src/types/narrow.rs

@@ -5,6 +5,7 @@ use crate::semantic_index::place_table;
 use crate::semantic_index::predicate::{
     CallableAndCallExpr, PatternPredicate, PatternPredicateKind, Predicate, PredicateNode,
 };
+use crate::types::enums::{enum_member_literals, enum_metadata};
 use crate::types::function::KnownFunction;
 use crate::types::infer::infer_same_file_expression_type;
 use crate::types::{
@@ -559,6 +560,16 @@ impl<'db, 'ast> NarrowingConstraintsBuilder<'db, 'ast> {
                                 .map(|ty| filter_to_cannot_be_equal(db, ty, rhs_ty)),
                         )
                     }
+                    // Treat enums as a union of their members.
+                    Type::NominalInstance(instance)
+                        if enum_metadata(db, instance.class.class_literal(db).0).is_some() =>
+                    {
+                        UnionType::from_elements(
+                            db,
+                            enum_member_literals(db, instance.class.class_literal(db).0, None)
+                                .map(|ty| filter_to_cannot_be_equal(db, ty, rhs_ty)),
+                        )
+                    }
                     _ => {
                         if ty.is_single_valued(db) && !could_compare_equal(db, ty, rhs_ty) {
                             ty