fix quick fix for Enum #3266

pyrefly/lib/alt/answers_solver.rs

@@ -76,6 +76,7 @@ use crate::error::collector::ErrorCollector;
 use crate::error::context::ErrorInfo;
 use crate::error::context::TypeCheckContext;
 use crate::error::context::TypeCheckKind;
+use crate::error::error::ErrorQuickFix;
 use crate::error::style::ErrorStyle;
 use crate::export::exports::LookupExport;
 use crate::module::module_info::ModuleInfo;
@@ -3058,11 +3059,16 @@ impl<'a, Ans: LookupAnswer> AnswersSolver<'a, Ans> {
         match subset_result {
             Ok(()) => true,
             Err(error) => {
-                let note = self
-                    .suggest_enum_member_for_value(got, want)
+                let enum_member_suggestion = self.suggest_enum_member_for_value(got, want);
+                let note = enum_member_suggestion
+                    .as_ref()
                     .map(|s| format!("Did you mean `{s}`?"));
+                let quick_fixes = enum_member_suggestion
+                    .map(|replacement| ErrorQuickFix::ReplaceWithEnumMember { replacement })
+                    .into_iter()
+                    .collect();
                 self.solver()
-                    .error(got, want, errors, loc, tcc, error, note);
+                    .error(got, want, errors, loc, tcc, error, note, quick_fixes);
                 false
             }
         }

pyrefly/lib/alt/class/typed_dict.rs

@@ -185,6 +185,7 @@ impl<'a, Ans: LookupAnswer> AnswersSolver<'a, Ans> {
                         tcc,
                         subset_error,
                         None,
+                        Vec::new(),
                     );
                 }
             }

pyrefly/lib/error/collector.rs

@@ -22,6 +22,7 @@ use crate::config::error::ErrorConfig;
 use crate::config::error_kind::Severity;
 use crate::error::context::ErrorInfo;
 use crate::error::error::Error;
+use crate::error::error::ErrorQuickFix;
 use crate::error::style::ErrorStyle;
 use crate::module::module_info::ModuleInfo;
 use crate::state::errors::find_containing_range;
@@ -131,34 +132,28 @@ impl ErrorCollector {
         }
     }
 
-    pub fn add(&self, range: TextRange, info: ErrorInfo, mut msg: Vec1<String>) {
-        if self.style == ErrorStyle::Never {
-            return;
-        }
-        let (kind, annotations) = match info {
-            ErrorInfo::Context(ctx) => {
-                let ctx = ctx();
-                let kind = ctx.as_error_kind();
-                let annotations = ctx.annotations();
-                msg.insert(0, ctx.format());
-                (kind, annotations)
-            }
-            ErrorInfo::Kind(kind) => (kind, Vec::new()),
-        };
-        let mut err = Error::new(self.module_info.dupe(), range, msg, kind);
-        for (range, label) in annotations {
-            err = err.with_annotation(range, label);
-        }
-        self.errors.lock().push(err);
+    pub fn add(&self, range: TextRange, info: ErrorInfo, msg: Vec1<String>) {
+        self.add_with_annotations_and_quick_fixes(range, info, msg, Vec::new(), Vec::new());
     }
 
     /// Add an error with secondary annotations for richer diagnostics.
     pub fn add_with_annotations(
+        &self,
+        range: TextRange,
+        info: ErrorInfo,
+        msg: Vec1<String>,
+        annotations: Vec<(TextRange, String)>,
+    ) {
+        self.add_with_annotations_and_quick_fixes(range, info, msg, annotations, Vec::new());
+    }
+
+    pub fn add_with_annotations_and_quick_fixes(
         &self,
         range: TextRange,
         info: ErrorInfo,
         mut msg: Vec1<String>,
         annotations: Vec<(TextRange, String)>,
+        quick_fixes: Vec<ErrorQuickFix>,
     ) {
         if self.style == ErrorStyle::Never {
             return;
@@ -177,6 +172,9 @@ impl ErrorCollector {
         for (range, label) in ctx_annotations.into_iter().chain(annotations) {
             err = err.with_annotation(range, label);
         }
+        for quick_fix in quick_fixes {
+            err = err.with_quick_fix(quick_fix);
+        }
         self.errors.lock().push(err);
     }
 

pyrefly/lib/error/error.rs

@@ -45,6 +45,11 @@ pub struct SecondaryAnnotation {
     pub label: Box<str>,
 }
 
+#[derive(Debug, Clone, PartialEq, Eq, Hash)]
+pub enum ErrorQuickFix {
+    ReplaceWithEnumMember { replacement: String },
+}
+
 #[derive(Debug, Clone, PartialEq, Eq, Hash)]
 pub struct Error {
     module: Module,
@@ -59,6 +64,8 @@ pub struct Error {
     msg_details: Option<Box<str>>,
     /// Additional labeled spans in the same file for richer diagnostics.
     secondary_annotations: Vec<SecondaryAnnotation>,
+    /// Structured fixes that can be exposed by editor integrations.
+    quick_fixes: Vec<ErrorQuickFix>,
 }
 
 impl Ranged for Error {
@@ -353,6 +360,7 @@ impl Error {
             msg_header,
             msg_details,
             secondary_annotations: Vec::new(),
+            quick_fixes: Vec::new(),
         }
     }
 
@@ -366,6 +374,11 @@ impl Error {
         self
     }
 
+    pub fn with_quick_fix(mut self, quick_fix: ErrorQuickFix) -> Self {
+        self.quick_fixes.push(quick_fix);
+        self
+    }
+
     pub fn display_range(&self) -> &DisplayRange {
         &self.display_range
     }
@@ -416,6 +429,10 @@ impl Error {
     pub fn secondary_annotations(&self) -> &[SecondaryAnnotation] {
         &self.secondary_annotations
     }
+
+    pub fn quick_fixes(&self) -> &[ErrorQuickFix] {
+        &self.quick_fixes
+    }
 }
 
 #[cfg(test)]

pyrefly/lib/solver/solver.rs

@@ -52,6 +52,7 @@ use crate::error::collector::ErrorCollector;
 use crate::error::context::ErrorInfo;
 use crate::error::context::TypeCheckContext;
 use crate::error::context::TypeCheckKind;
+use crate::error::error::ErrorQuickFix;
 use crate::solver::type_order::TypeOrder;
 use crate::types::callable::Callable;
 use crate::types::callable::Function;
@@ -1993,6 +1994,7 @@ impl Solver {
         tcc: &dyn Fn() -> TypeCheckContext,
         subset_error: SubsetError,
         note: Option<String>,
+        quick_fixes: Vec<ErrorQuickFix>,
     ) {
         let tcc = tcc();
         let msg = tcc.kind.format_error(
@@ -2010,19 +2012,21 @@ impl Solver {
         let extra_annotations = tcc.annotations;
         match tcc.context {
             Some(ctx) => {
-                errors.add_with_annotations(
+                errors.add_with_annotations_and_quick_fixes(
                     loc,
                     ErrorInfo::Context(&|| ctx.clone()),
                     msg_lines,
                     extra_annotations,
+                    quick_fixes,
                 );
             }
             None => {
-                errors.add_with_annotations(
+                errors.add_with_annotations_and_quick_fixes(
                     loc,
                     ErrorInfo::Kind(tcc.kind.as_error_kind()),
                     msg_lines,
                     extra_annotations,
+                    quick_fixes,
                 );
             }
         }

pyrefly/lib/state/lsp.rs

@@ -2347,6 +2347,18 @@ impl<'a> Transaction<'a> {
         let mut other_action_keys: HashSet<(String, TextRange, String)> = HashSet::new();
         for error in errors {
             let error_range = error.range();
+            if error_range.contains_range(range)
+                && let Some(action) = quick_fixes::enum_member::replace_with_enum_member_code_action(
+                    &module_info,
+                    &ast,
+                    &error,
+                )
+            {
+                let key = (action.0.clone(), action.2, action.3.clone());
+                if other_action_keys.insert(key) {
+                    other_actions.push(action);
+                }
+            }
             if error_range.contains_range(range)
                 && let Some(action) = quick_fixes::pyrefly_ignore::add_pyrefly_ignore_code_action(
                     &module_info,

pyrefly/lib/state/lsp/quick_fixes/enum_member.rs

@@ -0,0 +1,49 @@
+/*
+ * Copyright (c) Meta Platforms, Inc. and affiliates.
+ *
+ * This source code is licensed under the MIT license found in the
+ * LICENSE file in the root directory of this source tree.
+ */
+
+use dupe::Dupe;
+use pyrefly_python::ast::Ast;
+use pyrefly_python::module::Module;
+use ruff_python_ast::AnyNodeRef;
+use ruff_python_ast::ModModule;
+use ruff_text_size::Ranged;
+use ruff_text_size::TextRange;
+
+use crate::ModuleInfo;
+use crate::error::error::Error;
+use crate::error::error::ErrorQuickFix;
+
+pub(crate) fn replace_with_enum_member_code_action(
+    module_info: &ModuleInfo,
+    ast: &ModModule,
+    error: &Error,
+) -> Option<(String, Module, TextRange, String)> {
+    let replacement = enum_member_replacement(error)?;
+    let literal_range = enclosing_string_literal_range(ast, error.range())?;
+    Some((
+        format!("Replace with `{replacement}`"),
+        module_info.dupe(),
+        literal_range,
+        replacement.to_owned(),
+    ))
+}
+
+fn enum_member_replacement(error: &Error) -> Option<&str> {
+    let ErrorQuickFix::ReplaceWithEnumMember { replacement } = error.quick_fixes().first()?;
+    Some(replacement.as_str())
+}
+
+fn enclosing_string_literal_range(ast: &ModModule, error_range: TextRange) -> Option<TextRange> {
+    for node in Ast::locate_node(ast, error_range.start()) {
+        if let AnyNodeRef::ExprStringLiteral(literal) = node
+            && literal.range().contains_range(error_range)
+        {
+            return Some(literal.range());
+        }
+    }
+    None
+}

pyrefly/lib/state/lsp/quick_fixes/mod.rs

@@ -6,6 +6,7 @@
  */
 
 pub(crate) mod convert_star_import;
+pub(crate) mod enum_member;
 pub(crate) mod extract_field;
 pub(crate) mod extract_function;
 mod extract_shared;

pyrefly/lib/test/lsp/code_actions.rs

@@ -912,6 +912,35 @@ x: int = "hello"
     );
 }
 
+#[test]
+fn quickfix_replace_string_literal_with_enum_member() {
+    let report = get_batched_lsp_operations_report_allow_error(
+        &[(
+            "main",
+            r#"from enum import Enum
+
+class AccountStatus(Enum):
+    ACTIVE = "active"
+
+def takes_status(status: AccountStatus) -> None:
+    pass
+
+takes_status("active")
+#             ^
+"#,
+        )],
+        get_test_report,
+    );
+    assert!(
+        report.contains("# Title: Replace with `AccountStatus.ACTIVE`"),
+        "{report}"
+    );
+    assert!(
+        report.contains("takes_status(AccountStatus.ACTIVE)"),
+        "{report}"
+    );
+}
+
 #[test]
 fn quickfix_add_pyrefly_ignore_code_with_existing_comment() {
     let report = get_batched_lsp_operations_report_allow_error(