Fix enum Self type in match statements and circular dependency

crates/pyrefly_types/src/class.rs

@@ -86,6 +86,10 @@ pub struct ClassFieldProperties {
     is_initialized_on_class: bool,
     // The field is defined in the class body (not in a method via self.x = ...)
     is_defined_in_class_body: bool,
+    // Whether this field could potentially be an enum member (assigned or defined in the class
+    // body, not a method or nested class). Used to skip fields early in enum member enumeration
+    // to avoid circular dependencies when resolving field types.
+    could_be_enum_member: bool,
     range: TextRange,
     // The range of the docstring following this field, if present
     docstring_range: Option<TextRange>,
@@ -108,11 +112,17 @@ impl ClassFieldProperties {
             is_annotated,
             is_initialized_on_class: has_default_value,
             is_defined_in_class_body,
+            could_be_enum_member: false,
             range,
             docstring_range,
         }
     }
 
+    pub fn with_could_be_enum_member(mut self, could_be_enum_member: bool) -> Self {
+        self.could_be_enum_member = could_be_enum_member;
+        self
+    }
+
     pub fn is_initialized_on_class(&self) -> bool {
         self.is_initialized_on_class
     }
@@ -289,6 +299,16 @@ impl Class {
         self.0.fields.get(name)?.docstring_range
     }
 
+    /// Whether a field could potentially be an enum member based on its definition form.
+    /// This is a cheap check that avoids resolving field types, used to skip non-candidate
+    /// fields early in enum member enumeration and break circular dependencies.
+    pub fn could_field_be_enum_member(&self, name: &Name) -> bool {
+        self.0
+            .fields
+            .get(name)
+            .is_some_and(|prop| prop.could_be_enum_member)
+    }
+
     pub fn has_qname(&self, module: &str, parent: &NestingContext, name: &str) -> bool {
         self.0.qname.module_name().as_str() == module
             && self.0.qname.parent() == parent

pyrefly/lib/alt/class/enums.rs

@@ -50,6 +50,7 @@ impl<'a, Ans: LookupAnswer> AnswersSolver<'a, Ans> {
 
     pub fn get_enum_members(&self, cls: &Class) -> SmallSet<Lit> {
         cls.fields()
+            .filter(|f| cls.could_field_be_enum_member(f))
             .filter_map(|f| self.get_enum_member(cls, f))
             .collect()
     }

pyrefly/lib/alt/narrow.rs

@@ -958,9 +958,11 @@ impl<'a, Ans: LookupAnswer> AnswersSolver<'a, Ans> {
                                 {
                                     result = Lit::Bool(!b).to_implicit_type();
                                 }
-                                (Type::ClassType(left_cls), Type::Literal(right))
-                                    if let Lit::Enum(right) = &right.value
-                                        && left_cls == &right.class =>
+                                (
+                                    Type::ClassType(left_cls) | Type::SelfType(left_cls),
+                                    Type::Literal(right),
+                                ) if let Lit::Enum(right) = &right.value
+                                    && left_cls == &right.class =>
                                 {
                                     result = self.subtract_enum_member(left_cls, &right.member);
                                 }
@@ -1021,9 +1023,11 @@ impl<'a, Ans: LookupAnswer> AnswersSolver<'a, Ans> {
                         {
                             Lit::Bool(!b).to_implicit_type()
                         }
-                        (Type::ClassType(left_cls), Type::Literal(right))
-                            if let Lit::Enum(right) = &right.value
-                                && left_cls == &right.class =>
+                        (
+                            Type::ClassType(left_cls) | Type::SelfType(left_cls),
+                            Type::Literal(right),
+                        ) if let Lit::Enum(right) = &right.value
+                            && left_cls == &right.class =>
                         {
                             self.subtract_enum_member(left_cls, &right.member)
                         }
@@ -1212,9 +1216,11 @@ impl<'a, Ans: LookupAnswer> AnswersSolver<'a, Ans> {
                         {
                             Lit::Bool(!b).to_implicit_type()
                         }
-                        (Type::ClassType(left_cls), Type::Literal(right))
-                            if let Lit::Enum(right) = &right.value
-                                && left_cls == &right.class =>
+                        (
+                            Type::ClassType(left_cls) | Type::SelfType(left_cls),
+                            Type::Literal(right),
+                        ) if let Lit::Enum(right) = &right.value
+                            && left_cls == &right.class =>
                         {
                             self.subtract_enum_member(left_cls, &right.member)
                         }

pyrefly/lib/alt/overload.rs

@@ -163,7 +163,7 @@ impl<'a, Ans: LookupAnswer> ArgsExpander<'a, Ans> {
                     Lit::Bool(false).to_implicit_type(),
                 ]
             }
-            Type::ClassType(cls)
+            Type::ClassType(cls) | Type::SelfType(cls)
                 if self
                     .solver
                     .get_metadata_for_class(cls.class_object())

pyrefly/lib/binding/class.rs

@@ -405,6 +405,15 @@ impl<'a> BindingsBuilder<'a> {
                     }
                     _ => (true, false, true),
                 };
+            // Only AssignedInBody and DefinedWithoutAssign can be enum members
+            // (matching the check in is_valid_enum_member). This flag lets
+            // get_enum_members skip methods and nested classes without resolving
+            // their types, breaking circular dependencies.
+            let could_be_enum_member = matches!(
+                &definition,
+                ClassFieldDefinition::AssignedInBody { .. }
+                    | ClassFieldDefinition::DefinedWithoutAssign { .. }
+            );
 
             let docstring_range = field_docstrings.get(&range).copied();
 
@@ -416,7 +425,8 @@ impl<'a> BindingsBuilder<'a> {
                     is_defined_in_class_body,
                     range,
                     docstring_range,
-                ),
+                )
+                .with_could_be_enum_member(could_be_enum_member),
             );
             let key_field = KeyClassField(class_indices.def_index, name.clone().into_key());
             let binding = BindingClassField {

pyrefly/lib/solver/subset.rs

@@ -1732,7 +1732,8 @@ impl<'a, Ans: LookupAnswer> Subset<'a, Ans> {
             (Type::Literal(lit), Type::LiteralString(_)) => {
                 ok_or(lit.value.is_string(), SubsetError::Other)
             }
-            (Type::Literal(lit), t @ Type::ClassType(_)) => self.is_subset_eq(
+            (Type::Literal(lit), t @ Type::ClassType(_))
+            | (Type::Literal(lit), t @ Type::SelfType(_)) => self.is_subset_eq(
                 &self.solver.heap.mk_class_type(
                     lit.value
                         .general_class_type(self.type_order.stdlib())

pyrefly/lib/test/pattern_match.rs

@@ -636,6 +636,7 @@ testcase!(
     test_exhaustiveness_in_enum_method,
     r#"
 from enum import Enum
+from typing import reveal_type
 
 class E(Enum):
     X = 1
@@ -644,8 +645,10 @@ class E(Enum):
     def f_exhaustive(self) -> str:
         match self:
             case E.X:
+                reveal_type(self)  # E: Literal[E.X]
                 return "X"
             case E.Y:
+                reveal_type(self)  # E: Literal[E.Y]
                 return "Y"
 
     def f_nonexhaustive(self) -> str:  # E: missing an explicit `return`
@@ -655,6 +658,39 @@ class E(Enum):
     "#,
 );
 
+// Regression test for https://github.com/facebook/pyrefly/issues/2604
+// Calling an enum method that uses `match self` from outside should return
+// the declared return type, not Any.
+testcase!(
+    test_enum_method_return_type,
+    r#"
+from enum import Enum
+from typing import assert_type
+
+class A(Enum):
+    FOO = 1
+    BAR = 2
+
+    def method_with_match(self) -> int:
+        match self:
+            case A.FOO: return 1
+            case A.BAR: return 2
+
+    def method_with_if(self) -> int:
+        if self == A.FOO:
+            return 1
+        return 2
+
+    def method_no_ref(self) -> int:
+        return 1
+
+def test(a: A) -> None:
+    assert_type(a.method_no_ref(), int)
+    assert_type(a.method_with_if(), int)
+    assert_type(a.method_with_match(), int)
+    "#,
+);
+
 testcase!(
     test_match_mapping_after_none,
     r#"