[ty] equality narrowing on enums that don't override __eq__ or __ne__ (#20285)

Add equality narrowing for enums, if they don't override `__eq__` or `__ne__` in an unsafe way.

Follow-up to PR #20164

Fixes astral-sh/ty#939

Author: Renkai

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

@@ -168,11 +168,9 @@ class Color(Enum):
 
 def _(x: Color):
     if x in (Color.RED, Color.GREEN):
-        # TODO should be `Literal[Color.RED, Color.GREEN]`
-        reveal_type(x)  # revealed: Color
+        reveal_type(x)  # revealed: Literal[Color.RED, Color.GREEN]
     else:
-        # TODO should be `Literal[Color.BLUE]`
-        reveal_type(x)  # revealed: Color
+        reveal_type(x)  # revealed: Literal[Color.BLUE]
 ```
 
 ## Union with enum and `int`
@@ -187,11 +185,9 @@ class Status(Enum):
 
 def test(x: Status | int):
     if x in (Status.PENDING, Status.APPROVED):
-        # TODO should be `Literal[Status.PENDING, Status.APPROVED] | int`
         # int is included because custom __eq__ methods could make
         # an int equal to Status.PENDING or Status.APPROVED, so we can't eliminate it
-        reveal_type(x)  # revealed: Status | int
+        reveal_type(x)  # revealed: Literal[Status.PENDING, Status.APPROVED] | int
     else:
-        # TODO should be `Literal[Status.REJECTED] | int`
-        reveal_type(x)  # revealed: Status | int
+        reveal_type(x)  # revealed: Literal[Status.REJECTED] | int
 ```

crates/ty_python_semantic/src/types.rs

@@ -746,6 +746,35 @@ impl<'db> Type<'db> {
             .is_some_and(|instance| instance.class(db).is_known(db, KnownClass::Bool))
     }
 
+    fn is_enum(&self, db: &'db dyn Db) -> bool {
+        self.into_nominal_instance().is_some_and(|instance| {
+            crate::types::enums::enum_metadata(db, instance.class(db).class_literal(db).0).is_some()
+        })
+    }
+
+    /// Return true if this type overrides __eq__ or __ne__ methods
+    fn overrides_equality(&self, db: &'db dyn Db) -> bool {
+        let check_dunder = |dunder_name, allowed_return_value| {
+            // Note that we do explicitly exclude dunder methods on `object`, `int` and `str` here.
+            // The reason for this is that we know that these dunder methods behave in a predictable way.
+            // Only custom dunder methods need to be examined here, as they might break single-valuedness
+            // by always returning `False`, for example.
+            let call_result = self.try_call_dunder_with_policy(
+                db,
+                dunder_name,
+                &mut CallArguments::positional([Type::unknown()]),
+                MemberLookupPolicy::MRO_NO_OBJECT_FALLBACK
+                    | MemberLookupPolicy::MRO_NO_INT_OR_STR_LOOKUP,
+            );
+            let call_result = call_result.as_ref();
+            call_result.is_ok_and(|bindings| {
+                bindings.return_type(db) == Type::BooleanLiteral(allowed_return_value)
+            }) || call_result.is_err_and(|err| matches!(err, CallDunderError::MethodNotAvailable))
+        };
+
+        !(check_dunder("__eq__", true) && check_dunder("__ne__", false))
+    }
+
     pub(crate) fn is_notimplemented(&self, db: &'db dyn Db) -> bool {
         self.into_nominal_instance().is_some_and(|instance| {
             instance
@@ -980,22 +1009,28 @@ impl<'db> Type<'db> {
 
     pub(crate) fn is_union_of_single_valued(&self, db: &'db dyn Db) -> bool {
         self.into_union().is_some_and(|union| {
-            union
-                .elements(db)
-                .iter()
-                .all(|ty| ty.is_single_valued(db) || ty.is_bool(db) || ty.is_literal_string())
+            union.elements(db).iter().all(|ty| {
+                ty.is_single_valued(db)
+                    || ty.is_bool(db)
+                    || ty.is_literal_string()
+                    || (ty.is_enum(db) && !ty.overrides_equality(db))
+            })
         }) || self.is_bool(db)
             || self.is_literal_string()
+            || (self.is_enum(db) && !self.overrides_equality(db))
     }
 
     pub(crate) fn is_union_with_single_valued(&self, db: &'db dyn Db) -> bool {
         self.into_union().is_some_and(|union| {
-            union
-                .elements(db)
-                .iter()
-                .any(|ty| ty.is_single_valued(db) || ty.is_bool(db) || ty.is_literal_string())
+            union.elements(db).iter().any(|ty| {
+                ty.is_single_valued(db)
+                    || ty.is_bool(db)
+                    || ty.is_literal_string()
+                    || (ty.is_enum(db) && !ty.overrides_equality(db))
+            })
         }) || self.is_bool(db)
             || self.is_literal_string()
+            || (self.is_enum(db) && !self.overrides_equality(db))
     }
 
     pub(crate) fn into_string_literal(self) -> Option<StringLiteralType<'db>> {
@@ -2574,28 +2609,7 @@ impl<'db> Type<'db> {
             | Type::SpecialForm(..)
             | Type::KnownInstance(..) => true,
 
-            Type::EnumLiteral(_) => {
-                let check_dunder = |dunder_name, allowed_return_value| {
-                    // Note that we do explicitly exclude dunder methods on `object`, `int` and `str` here.
-                    // The reason for this is that we know that these dunder methods behave in a predictable way.
-                    // Only custom dunder methods need to be examined here, as they might break single-valuedness
-                    // by always returning `False`, for example.
-                    let call_result = self.try_call_dunder_with_policy(
-                        db,
-                        dunder_name,
-                        &mut CallArguments::positional([Type::unknown()]),
-                        MemberLookupPolicy::MRO_NO_OBJECT_FALLBACK
-                            | MemberLookupPolicy::MRO_NO_INT_OR_STR_LOOKUP,
-                    );
-                    let call_result = call_result.as_ref();
-                    call_result.is_ok_and(|bindings| {
-                        bindings.return_type(db) == Type::BooleanLiteral(allowed_return_value)
-                    }) || call_result
-                        .is_err_and(|err| matches!(err, CallDunderError::MethodNotAvailable))
-                };
-
-                check_dunder("__eq__", true) && check_dunder("__ne__", false)
-            }
+            Type::EnumLiteral(_) => !self.overrides_equality(db),
 
             Type::ProtocolInstance(..) => {
                 // See comment in the `Type::ProtocolInstance` branch for `Type::is_singleton`.

crates/ty_python_semantic/src/types/narrow.rs

@@ -640,6 +640,7 @@ impl<'db, 'ast> NarrowingConstraintsBuilder<'db, 'ast> {
                     if !element.is_single_valued(self.db)
                         && !element.is_literal_string()
                         && !element.is_bool(self.db)
+                        && (!element.is_enum(self.db) || element.overrides_equality(self.db))
                     {
                         builder = builder.add(*element);
                     }
@@ -675,6 +676,7 @@ impl<'db, 'ast> NarrowingConstraintsBuilder<'db, 'ast> {
                     if element.is_single_valued(self.db)
                         || element.is_literal_string()
                         || element.is_bool(self.db)
+                        || (element.is_enum(self.db) && !element.overrides_equality(self.db))
                     {
                         single_builder = single_builder.add(*element);
                     } else {