From dae280bb9889f8a985fd78ae59f015c7e97c2abe Mon Sep 17 00:00:00 2001 From: Asuka Minato Date: Thu, 30 Apr 2026 19:16:16 +0900 Subject: [PATCH 1/2] fix --- pyrefly/lib/state/lsp.rs | 12 ++++ .../lib/state/lsp/quick_fixes/enum_member.rs | 70 +++++++++++++++++++ pyrefly/lib/state/lsp/quick_fixes/mod.rs | 1 + pyrefly/lib/test/lsp/code_actions.rs | 29 ++++++++ 4 files changed, 112 insertions(+) create mode 100644 pyrefly/lib/state/lsp/quick_fixes/enum_member.rs diff --git a/pyrefly/lib/state/lsp.rs b/pyrefly/lib/state/lsp.rs index b9c156729d..bdd9813a19 100644 --- a/pyrefly/lib/state/lsp.rs +++ b/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, diff --git a/pyrefly/lib/state/lsp/quick_fixes/enum_member.rs b/pyrefly/lib/state/lsp/quick_fixes/enum_member.rs new file mode 100644 index 0000000000..5265ce5388 --- /dev/null +++ b/pyrefly/lib/state/lsp/quick_fixes/enum_member.rs @@ -0,0 +1,70 @@ +/* + * 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; + +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_suggestion(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_suggestion(error: &Error) -> Option<&str> { + let details = error.msg_details()?; + for line in details.lines() { + let Some(suggestion) = line.trim().strip_prefix("Did you mean `") else { + continue; + }; + let Some(suggestion) = suggestion.strip_suffix("`?") else { + continue; + }; + if is_qualified_identifier(suggestion) { + return Some(suggestion); + } + } + None +} + +fn is_qualified_identifier(s: &str) -> bool { + s.split('.').count() >= 2 + && s.split('.').all(|part| { + let mut chars = part.chars(); + chars + .next() + .is_some_and(|c| c == '_' || c.is_ascii_alphabetic()) + && chars.all(|c| c == '_' || c.is_ascii_alphanumeric()) + }) +} + +fn enclosing_string_literal_range(ast: &ModModule, error_range: TextRange) -> Option { + 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 +} diff --git a/pyrefly/lib/state/lsp/quick_fixes/mod.rs b/pyrefly/lib/state/lsp/quick_fixes/mod.rs index 4c66d7ffb9..c416e64730 100644 --- a/pyrefly/lib/state/lsp/quick_fixes/mod.rs +++ b/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; diff --git a/pyrefly/lib/test/lsp/code_actions.rs b/pyrefly/lib/test/lsp/code_actions.rs index 682a85d669..01bf1b73c5 100644 --- a/pyrefly/lib/test/lsp/code_actions.rs +++ b/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( From d2e97d4174c07b4877c784ec750f2e5cf416c31f Mon Sep 17 00:00:00 2001 From: Asuka Minato Date: Fri, 1 May 2026 12:49:41 +0900 Subject: [PATCH 2/2] comment --- pyrefly/lib/alt/answers_solver.rs | 12 +++++-- pyrefly/lib/alt/class/typed_dict.rs | 1 + pyrefly/lib/error/collector.rs | 36 +++++++++---------- pyrefly/lib/error/error.rs | 17 +++++++++ pyrefly/lib/solver/solver.rs | 8 +++-- .../lib/state/lsp/quick_fixes/enum_member.rs | 31 +++------------- 6 files changed, 55 insertions(+), 50 deletions(-) diff --git a/pyrefly/lib/alt/answers_solver.rs b/pyrefly/lib/alt/answers_solver.rs index 241fb1aeb3..3c6d92d6cd 100644 --- a/pyrefly/lib/alt/answers_solver.rs +++ b/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 } } diff --git a/pyrefly/lib/alt/class/typed_dict.rs b/pyrefly/lib/alt/class/typed_dict.rs index 46edb824c9..40404768cc 100644 --- a/pyrefly/lib/alt/class/typed_dict.rs +++ b/pyrefly/lib/alt/class/typed_dict.rs @@ -185,6 +185,7 @@ impl<'a, Ans: LookupAnswer> AnswersSolver<'a, Ans> { tcc, subset_error, None, + Vec::new(), ); } } diff --git a/pyrefly/lib/error/collector.rs b/pyrefly/lib/error/collector.rs index 64f13895c2..49cb15ceff 100644 --- a/pyrefly/lib/error/collector.rs +++ b/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) { - 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) { + 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, + 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, annotations: Vec<(TextRange, String)>, + quick_fixes: Vec, ) { 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); } diff --git a/pyrefly/lib/error/error.rs b/pyrefly/lib/error/error.rs index d541c765c1..b38811a5a8 100644 --- a/pyrefly/lib/error/error.rs +++ b/pyrefly/lib/error/error.rs @@ -45,6 +45,11 @@ pub struct SecondaryAnnotation { pub label: Box, } +#[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>, /// Additional labeled spans in the same file for richer diagnostics. secondary_annotations: Vec, + /// Structured fixes that can be exposed by editor integrations. + quick_fixes: Vec, } 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)] diff --git a/pyrefly/lib/solver/solver.rs b/pyrefly/lib/solver/solver.rs index c5cf94080c..250fb833bf 100644 --- a/pyrefly/lib/solver/solver.rs +++ b/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, + quick_fixes: Vec, ) { 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, ); } } diff --git a/pyrefly/lib/state/lsp/quick_fixes/enum_member.rs b/pyrefly/lib/state/lsp/quick_fixes/enum_member.rs index 5265ce5388..d1d0c0c060 100644 --- a/pyrefly/lib/state/lsp/quick_fixes/enum_member.rs +++ b/pyrefly/lib/state/lsp/quick_fixes/enum_member.rs @@ -15,13 +15,14 @@ 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_suggestion(error)?; + let replacement = enum_member_replacement(error)?; let literal_range = enclosing_string_literal_range(ast, error.range())?; Some(( format!("Replace with `{replacement}`"), @@ -31,31 +32,9 @@ pub(crate) fn replace_with_enum_member_code_action( )) } -fn enum_member_suggestion(error: &Error) -> Option<&str> { - let details = error.msg_details()?; - for line in details.lines() { - let Some(suggestion) = line.trim().strip_prefix("Did you mean `") else { - continue; - }; - let Some(suggestion) = suggestion.strip_suffix("`?") else { - continue; - }; - if is_qualified_identifier(suggestion) { - return Some(suggestion); - } - } - None -} - -fn is_qualified_identifier(s: &str) -> bool { - s.split('.').count() >= 2 - && s.split('.').all(|part| { - let mut chars = part.chars(); - chars - .next() - .is_some_and(|c| c == '_' || c.is_ascii_alphabetic()) - && chars.all(|c| c == '_' || c.is_ascii_alphanumeric()) - }) +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 {