diff --git a/pyrefly/lib/lsp/non_wasm/server.rs b/pyrefly/lib/lsp/non_wasm/server.rs index 1e9eee5457..241a9321fd 100644 --- a/pyrefly/lib/lsp/non_wasm/server.rs +++ b/pyrefly/lib/lsp/non_wasm/server.rs @@ -5395,6 +5395,10 @@ impl Server { "invert_boolean", transaction.invert_boolean_code_actions(&handle, range) ); + timed_refactor_action!( + "if_to_match", + transaction.if_to_match_code_actions(&handle, range) + ); timed_refactor_action!( "extract_function", transaction.extract_function_code_actions(&handle, range) diff --git a/pyrefly/lib/state/lsp.rs b/pyrefly/lib/state/lsp.rs index 47531097a6..4befbe3428 100644 --- a/pyrefly/lib/state/lsp.rs +++ b/pyrefly/lib/state/lsp.rs @@ -4059,6 +4059,14 @@ impl<'a> Transaction<'a> { quick_fixes::invert_boolean::invert_boolean_code_actions(self, handle, selection) } + pub fn if_to_match_code_actions( + &self, + handle: &Handle, + selection: TextRange, + ) -> Option> { + quick_fixes::if_to_match::if_to_match_code_actions(self, handle, selection) + } + pub fn extract_superclass_code_actions( &self, handle: &Handle, diff --git a/pyrefly/lib/state/lsp/quick_fixes.rs b/pyrefly/lib/state/lsp/quick_fixes.rs index 2aeed6d263..ccf2aed13f 100644 --- a/pyrefly/lib/state/lsp/quick_fixes.rs +++ b/pyrefly/lib/state/lsp/quick_fixes.rs @@ -17,6 +17,7 @@ mod extract_shared; pub(crate) mod extract_superclass; pub(crate) mod extract_variable; pub(crate) mod generate_code; +pub(crate) mod if_to_match; pub(crate) mod inline_method; pub(crate) mod inline_parameter; pub(crate) mod inline_variable; diff --git a/pyrefly/lib/state/lsp/quick_fixes/if_to_match.rs b/pyrefly/lib/state/lsp/quick_fixes/if_to_match.rs new file mode 100644 index 0000000000..d3fd8d9f9f --- /dev/null +++ b/pyrefly/lib/state/lsp/quick_fixes/if_to_match.rs @@ -0,0 +1,244 @@ +/* + * 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 lsp_types::CodeActionKind; +use pyrefly_build::handle::Handle; +use pyrefly_python::ast::Ast; +use ruff_python_ast::AnyNodeRef; +use ruff_python_ast::CmpOp; +use ruff_python_ast::Expr; +use ruff_python_ast::ModModule; +use ruff_python_ast::Stmt; +use ruff_python_ast::StmtIf; +use ruff_text_size::Ranged; +use ruff_text_size::TextRange; +use ruff_text_size::TextSize; + +use super::extract_shared::line_indent_and_start; +use super::extract_shared::selection_anchor; +use crate::state::lsp::LocalRefactorCodeAction; +use crate::state::lsp::Transaction; + +/// Builds an if/elif-to-match refactor for simple same-subject narrowing chains. +pub(crate) fn if_to_match_code_actions( + transaction: &Transaction<'_>, + handle: &Handle, + selection: TextRange, +) -> Option> { + let module_info = transaction.get_module_info(handle)?; + let ast = transaction.get_ast(handle)?; + let source = module_info.contents(); + let selection_point = selection_anchor(source, selection); + let if_stmt = find_if_at_selection(ast.as_ref(), selection_point)?; + let replacement = build_match_replacement(source, if_stmt)?; + let (_, line_start) = line_indent_and_start(source, if_stmt.range().start())?; + let replace_range = + TextRange::new(line_start, line_end_position(source, if_stmt.range().end())); + + Some(vec![LocalRefactorCodeAction { + title: "Convert if/elif chain to match".to_owned(), + edits: vec![(module_info.dupe(), replace_range, replacement)], + kind: CodeActionKind::RefactorRewrite, + }]) +} + +fn find_if_at_selection(ast: &ModModule, selection_point: TextSize) -> Option<&StmtIf> { + Ast::locate_node(ast, selection_point) + .into_iter() + .find_map(|node| match node { + AnyNodeRef::StmtIf(if_stmt) => Some(if_stmt), + _ => None, + }) +} + +fn build_match_replacement(source: &str, if_stmt: &StmtIf) -> Option { + let (base_indent, _) = line_indent_and_start(source, if_stmt.range().start())?; + let body_indent = branch_body_indent(source, &if_stmt.body, &base_indent)?; + let mut subject = None; + let mut cases = Vec::new(); + let if_end = line_end_position(source, if_stmt.range().end()); + let first_body_end = if_stmt + .elif_else_clauses + .first() + .and_then(|clause| line_start_position(source, clause.range.start())) + .unwrap_or(if_end); + cases.push(branch_case( + source, + &if_stmt.test, + TextRange::new( + line_end_position(source, if_stmt.test.range().end()), + first_body_end, + ), + &mut subject, + )?); + for (index, clause) in if_stmt.elif_else_clauses.iter().enumerate() { + let body_end = if_stmt + .elif_else_clauses + .get(index + 1) + .and_then(|next_clause| line_start_position(source, next_clause.range.start())) + .unwrap_or(if_end); + match &clause.test { + Some(test) => cases.push(branch_case( + source, + test, + TextRange::new(line_end_position(source, test.range().end()), body_end), + &mut subject, + )?), + None => cases.push(else_case( + source, + TextRange::new(line_end_position(source, clause.range.start()), body_end), + )?), + } + } + if cases.iter().filter(|case| case.pattern != "_").count() < 2 { + return None; + } + + let subject = subject?; + let case_indent = format!("{base_indent}{body_indent}"); + let mut replacement = format!("{base_indent}match {subject}:\n"); + for case in cases { + replacement.push_str(&case_indent); + replacement.push_str("case "); + replacement.push_str(&case.pattern); + replacement.push_str(":\n"); + replacement.push_str(&indent_branch_body(&case.body_text, &body_indent)); + } + Some(replacement) +} + +struct MatchCaseText { + pattern: String, + body_text: String, +} + +fn branch_case( + source: &str, + test: &Expr, + body_range: TextRange, + subject: &mut Option, +) -> Option { + let (candidate_subject, pattern) = case_subject_and_pattern(source, test)?; + match subject { + Some(subject) if subject != &candidate_subject => return None, + Some(_) => {} + None => *subject = Some(candidate_subject), + } + Some(MatchCaseText { + pattern, + body_text: source_text(source, body_range).to_owned(), + }) +} + +fn else_case(source: &str, body_range: TextRange) -> Option { + Some(MatchCaseText { + pattern: "_".to_owned(), + body_text: source_text(source, body_range).to_owned(), + }) +} + +fn case_subject_and_pattern(source: &str, test: &Expr) -> Option<(String, String)> { + if let Expr::Compare(compare) = test + && compare.ops.as_ref() == [CmpOp::Eq] + { + let subject = expr_text(source, compare.first_operand()); + let pattern = value_pattern(source, compare.comparators().first()?)?; + return Some((subject, pattern)); + } + if let Expr::Call(call) = test + && let Expr::Name(func) = call.func.as_ref() + && func.id.as_str() == "isinstance" + && call.arguments.args.len() == 2 + && call.arguments.keywords.is_empty() + { + let subject = expr_text(source, &call.arguments.args[0]); + let class_pattern = class_pattern(source, &call.arguments.args[1])?; + return Some((subject, class_pattern)); + } + None +} + +fn value_pattern(source: &str, value: &Expr) -> Option { + match value { + Expr::StringLiteral(_) + | Expr::BytesLiteral(_) + | Expr::NumberLiteral(_) + | Expr::BooleanLiteral(_) + | Expr::NoneLiteral(_) + | Expr::Attribute(_) => Some(expr_text(source, value)), + _ => None, + } +} + +fn class_pattern(source: &str, value: &Expr) -> Option { + match value { + Expr::Name(_) | Expr::Attribute(_) => Some(format!("{}()", expr_text(source, value))), + Expr::Tuple(tuple) if !tuple.elts.is_empty() => { + let patterns: Option> = tuple + .elts + .iter() + .map(|elt| match elt { + Expr::Name(_) | Expr::Attribute(_) => { + Some(format!("{}()", expr_text(source, elt))) + } + _ => None, + }) + .collect(); + Some(patterns?.join(" | ")) + } + _ => None, + } +} + +fn expr_text(source: &str, expr: &Expr) -> String { + source_text(source, expr.range()).trim().to_owned() +} + +fn branch_body_indent(source: &str, body: &[Stmt], base_indent: &str) -> Option { + let (indent, _) = line_indent_and_start(source, body.first()?.range().start())?; + Some( + indent + .strip_prefix(base_indent) + .filter(|indent| !indent.is_empty()) + .unwrap_or(" ") + .to_owned(), + ) +} + +fn indent_branch_body(body: &str, extra_indent: &str) -> String { + let mut result = String::new(); + for line in body.split_inclusive('\n') { + if line.trim().is_empty() { + result.push_str(line); + } else { + result.push_str(extra_indent); + result.push_str(line); + } + } + if !body.ends_with('\n') { + result.push('\n'); + } + result +} + +fn source_text(source: &str, range: TextRange) -> &str { + &source[range.start().to_usize().min(source.len())..range.end().to_usize().min(source.len())] +} + +fn line_end_position(source: &str, position: TextSize) -> TextSize { + let idx = position.to_usize().min(source.len()); + if let Some(offset) = source[idx..].find('\n') { + TextSize::try_from(idx + offset + 1).unwrap_or(position) + } else { + TextSize::try_from(source.len()).unwrap_or(position) + } +} + +fn line_start_position(source: &str, position: TextSize) -> Option { + line_indent_and_start(source, position).map(|(_, start)| start) +} diff --git a/pyrefly/lib/test/lsp/code_actions.rs b/pyrefly/lib/test/lsp/code_actions.rs index 2b83c8698a..506d7637a8 100644 --- a/pyrefly/lib/test/lsp/code_actions.rs +++ b/pyrefly/lib/test/lsp/code_actions.rs @@ -304,6 +304,43 @@ fn assert_no_invert_boolean_action_allow_errors(code: &str, selection: TextRange ); } +fn compute_if_to_match_actions( + code: &str, + selection: TextRange, +) -> ( + ModuleInfo, + Vec>, + Vec, +) { + let (handles, state) = + mk_multi_file_state_assert_no_errors(&[("main", code)], Require::Everything); + let handle = handles.get("main").unwrap(); + let transaction = state.transaction(); + let module_info = transaction.get_module_info(handle).unwrap(); + let actions = transaction + .if_to_match_code_actions(handle, selection) + .unwrap_or_default(); + let edit_sets: Vec> = + actions.iter().map(|action| action.edits.clone()).collect(); + let titles = actions.iter().map(|action| action.title.clone()).collect(); + (module_info, edit_sets, titles) +} + +fn apply_first_if_to_match_action(code: &str, selection: TextRange) -> Option { + let (module_info, actions, _) = compute_if_to_match_actions(code, selection); + let edits = actions.first()?; + Some(apply_refactor_edits_for_module(&module_info, edits)) +} + +fn assert_no_if_to_match_action(code: &str, selection: TextRange) { + let (_, actions, _) = compute_if_to_match_actions(code, selection); + assert!( + actions.is_empty(), + "expected no if-to-match actions, found {}", + actions.len() + ); +} + fn cursor_selection(code: &str) -> TextRange { let position = extract_cursors_for_test(code) .first() @@ -3577,6 +3614,176 @@ def foo(other_var): assert_eq!(expected.trim(), updated.trim()); } +#[test] +fn if_to_match_literal_chain() { + let code = r#" +x = "a" +if x == "a": + y = 1 +elif x == "b": + y = 2 +else: + y = 3 +"#; + let selection = find_nth_range(code, "if x", 1); + let updated = + apply_first_if_to_match_action(code, selection).expect("expected if-to-match action"); + let expected = r#" +x = "a" +match x: + case "a": + y = 1 + case "b": + y = 2 + case _: + y = 3 +"#; + assert_eq!(expected.trim(), updated.trim()); +} + +#[test] +fn if_to_match_isinstance_chain() { + let code = r#" +def f(x): + if isinstance(x, str): + return x + elif isinstance(x, int): + return str(x) +"#; + let selection = find_nth_range(code, "if isinstance", 1); + let updated = + apply_first_if_to_match_action(code, selection).expect("expected if-to-match action"); + let expected = r#" +def f(x): + match x: + case str(): + return x + case int(): + return str(x) +"#; + assert_eq!(expected.trim(), updated.trim()); +} + +#[test] +fn if_to_match_isinstance_tuple_chain() { + let code = r#" +class A: pass +class B: pass +class C: pass + +def f(x): + if isinstance(x, (A, B)): + return 1 + elif isinstance(x, C): + return 2 +"#; + let selection = find_nth_range(code, "if isinstance", 1); + let updated = + apply_first_if_to_match_action(code, selection).expect("expected if-to-match action"); + let expected = r#" +class A: pass +class B: pass +class C: pass + +def f(x): + match x: + case A() | B(): + return 1 + case C(): + return 2 +"#; + assert_eq!(expected.trim(), updated.trim()); +} + +#[test] +fn if_to_match_enum_member_chain() { + let code = r#" +from enum import Enum + +class Color(Enum): + RED = 1 + BLUE = 2 + +def f(color: Color): + if color == Color.RED: + return 1 + elif color == Color.BLUE: + return 2 +"#; + let selection = find_nth_range(code, "if color", 1); + let updated = + apply_first_if_to_match_action(code, selection).expect("expected if-to-match action"); + let expected = r#" +from enum import Enum + +class Color(Enum): + RED = 1 + BLUE = 2 + +def f(color: Color): + match color: + case Color.RED: + return 1 + case Color.BLUE: + return 2 +"#; + assert_eq!(expected.trim(), updated.trim()); +} + +#[test] +fn if_to_match_preserves_branch_comments() { + let code = r#" +x = "a" +if x == "a": + # keep this comment + y = 1 +elif x == "b": + y = 2 +"#; + let selection = find_nth_range(code, "if x", 1); + let updated = + apply_first_if_to_match_action(code, selection).expect("expected if-to-match action"); + let expected = r#" +x = "a" +match x: + case "a": + # keep this comment + y = 1 + case "b": + y = 2 +"#; + assert_eq!(expected.trim(), updated.trim()); +} + +#[test] +fn if_to_match_rejects_mixed_subjects() { + let code = r#" +x = "a" +z = "b" +if x == "a": + y = 1 +elif z == "b": + y = 2 +"#; + let selection = find_nth_range(code, "if x", 1); + assert_no_if_to_match_action(code, selection); +} + +#[test] +fn if_to_match_rejects_capture_pattern() { + let code = r#" +x = "a" +expected = "a" +other = "b" +if x == expected: + y = 1 +elif x == other: + y = 2 +"#; + let selection = find_nth_range(code, "if x", 1); + assert_no_if_to_match_action(code, selection); +} + #[test] fn pull_member_up_basic() { let code = r#"