diff --git a/pyrefly/lib/lsp/non_wasm/server.rs b/pyrefly/lib/lsp/non_wasm/server.rs index 1e9eee5457..d4f1b76533 100644 --- a/pyrefly/lib/lsp/non_wasm/server.rs +++ b/pyrefly/lib/lsp/non_wasm/server.rs @@ -5451,6 +5451,10 @@ impl Server { "convert_dict", transaction.convert_dict_code_actions(&handle, range) ); + timed_refactor_action!( + "unwrap_block", + transaction.unwrap_block_code_actions(&handle, range) + ); timed_refactor_action!( "pytest_fixture_type_annotation", transaction.pytest_fixture_type_annotation_code_actions( diff --git a/pyrefly/lib/state/lsp.rs b/pyrefly/lib/state/lsp.rs index 47531097a6..3c31463166 100644 --- a/pyrefly/lib/state/lsp.rs +++ b/pyrefly/lib/state/lsp.rs @@ -4190,6 +4190,14 @@ impl<'a> Transaction<'a> { quick_fixes::convert_dict::convert_dict_code_actions(self, handle, selection) } + pub fn unwrap_block_code_actions( + &self, + handle: &Handle, + selection: TextRange, + ) -> Option> { + quick_fixes::unwrap_block::unwrap_block_code_actions(self, handle, selection) + } + pub fn change_signature_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..e0d2f1bd6a 100644 --- a/pyrefly/lib/state/lsp/quick_fixes.rs +++ b/pyrefly/lib/state/lsp/quick_fixes.rs @@ -32,3 +32,4 @@ pub(crate) mod types; pub(crate) mod unnecessary_type_conversion; pub(crate) mod unused_ignore; pub(crate) mod unused_import; +pub(crate) mod unwrap_block; diff --git a/pyrefly/lib/state/lsp/quick_fixes/unwrap_block.rs b/pyrefly/lib/state/lsp/quick_fixes/unwrap_block.rs new file mode 100644 index 0000000000..11dcdca58c --- /dev/null +++ b/pyrefly/lib/state/lsp/quick_fixes/unwrap_block.rs @@ -0,0 +1,435 @@ +/* + * 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 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::code_at_range; +use super::extract_shared::line_end_position; +use super::extract_shared::line_indent_and_start; +use super::extract_shared::reindent_block; +use super::extract_shared::selection_anchor; +use super::extract_shared::statement_removal_range_from_range; +use super::types::LocalRefactorCodeAction; +use crate::state::lsp::Transaction; + +struct UnwrapBlockTarget { + outer_range: TextRange, + body_range: TextRange, + from_indent: String, + to_indent: String, +} + +#[derive(Clone, Copy, Debug)] +struct QuoteState { + delimiter: char, + triple: bool, +} + +/// Builds unwrap-block refactor actions for the block under the cursor. +pub(crate) fn unwrap_block_code_actions( + transaction: &Transaction<'_>, + handle: &Handle, + selection: TextRange, +) -> Option> { + let module_info = transaction.get_module_info(handle)?; + let source = module_info.contents(); + let ast = transaction.get_ast(handle)?; + let position = selection_anchor(source, selection); + let target = find_unwrap_block_target_in_body(&ast.body, source, position)?; + let removal_range = statement_removal_range_from_range(source, target.outer_range)?; + let replacement = reindent_block( + code_at_range(source, target.body_range)?, + &target.from_indent, + &target.to_indent, + ); + Some(vec![LocalRefactorCodeAction { + title: "Unwrap block".to_owned(), + edits: vec![(module_info.dupe(), removal_range, replacement)], + kind: CodeActionKind::RefactorRewrite, + }]) +} + +fn find_unwrap_block_target_in_body( + body: &[Stmt], + source: &str, + position: TextSize, +) -> Option { + for stmt in body { + if !range_contains_position(stmt.range(), position) { + continue; + } + if let Some(target) = find_unwrap_block_target_in_stmt(stmt, source, position) { + return Some(target); + } + } + None +} + +fn find_unwrap_block_target_in_stmt( + stmt: &Stmt, + source: &str, + position: TextSize, +) -> Option { + match stmt { + Stmt::FunctionDef(function_def) => { + find_unwrap_block_target_in_body(&function_def.body, source, position) + } + Stmt::ClassDef(class_def) => { + find_unwrap_block_target_in_body(&class_def.body, source, position) + } + Stmt::If(if_stmt) => find_unwrap_if_target(if_stmt, source, position), + Stmt::For(for_stmt) => { + if let Some(target) = find_unwrap_block_target_in_body(&for_stmt.body, source, position) + { + return Some(target); + } + if !for_stmt.orelse.is_empty() { + return None; + } + build_unwrap_block_target( + stmt.range(), + stmt.range().start(), + &for_stmt.body, + None, + source, + position, + ) + } + Stmt::While(while_stmt) => { + if let Some(target) = + find_unwrap_block_target_in_body(&while_stmt.body, source, position) + { + return Some(target); + } + if let Some(target) = + find_unwrap_block_target_in_body(&while_stmt.orelse, source, position) + { + return Some(target); + } + if !while_stmt.orelse.is_empty() { + return None; + } + build_unwrap_block_target( + stmt.range(), + stmt.range().start(), + &while_stmt.body, + None, + source, + position, + ) + } + Stmt::With(with_stmt) => { + if let Some(target) = + find_unwrap_block_target_in_body(&with_stmt.body, source, position) + { + return Some(target); + } + build_unwrap_block_target( + stmt.range(), + stmt.range().start(), + &with_stmt.body, + None, + source, + position, + ) + } + Stmt::Try(try_stmt) => { + let (to_indent, _) = line_indent_and_start(source, stmt.range().start())?; + if let Some(target) = find_unwrap_block_target_in_body(&try_stmt.body, source, position) + { + return Some(target); + } + for handler in &try_stmt.handlers { + if let Some(handler) = handler.as_except_handler() + && let Some(target) = + find_unwrap_block_target_in_body(&handler.body, source, position) + { + return Some(target); + } + } + if let Some(target) = + find_unwrap_block_target_in_body(&try_stmt.orelse, source, position) + { + return Some(target); + } + if let Some(target) = + find_unwrap_block_target_in_body(&try_stmt.finalbody, source, position) + { + return Some(target); + } + for (index, handler) in try_stmt.handlers.iter().enumerate() { + let Some(handler) = handler.as_except_handler() else { + continue; + }; + let next_clause_start = try_stmt + .handlers + .get(index + 1) + .and_then(|next| next.as_except_handler()) + .map(|next| next.range().start()) + .or_else(|| { + try_stmt + .orelse + .first() + .map(Ranged::range) + .map(|range| range.start()) + }) + .or_else(|| { + try_stmt + .finalbody + .first() + .map(Ranged::range) + .map(|range| range.start()) + }); + if let Some(target) = build_unwrap_block_target( + stmt.range(), + handler.range().start(), + &handler.body, + next_clause_start, + source, + position, + ) { + return Some(target); + } + } + let else_start = try_stmt.orelse.first().and_then(|stmt| { + find_try_clause_start(source, stmt.range().start(), &to_indent, "else") + }); + if let Some(else_start) = else_start + && let Some(target) = build_unwrap_block_target( + stmt.range(), + else_start, + &try_stmt.orelse, + try_stmt.finalbody.first().and_then(|stmt| { + find_try_clause_start(source, stmt.range().start(), &to_indent, "finally") + }), + source, + position, + ) + { + return Some(target); + } + let finally_start = try_stmt.finalbody.first().and_then(|stmt| { + find_try_clause_start(source, stmt.range().start(), &to_indent, "finally") + }); + build_unwrap_block_target( + stmt.range(), + finally_start?, + &try_stmt.finalbody, + None, + source, + position, + ) + } + Stmt::Match(match_stmt) => { + for case in &match_stmt.cases { + if let Some(target) = find_unwrap_block_target_in_body(&case.body, source, position) + { + return Some(target); + } + } + None + } + _ => None, + } +} + +fn find_unwrap_if_target( + if_stmt: &StmtIf, + source: &str, + position: TextSize, +) -> Option { + if let Some(target) = find_unwrap_block_target_in_body(&if_stmt.body, source, position) { + return Some(target); + } + for clause in &if_stmt.elif_else_clauses { + if let Some(target) = find_unwrap_block_target_in_body(&clause.body, source, position) { + return Some(target); + } + } + + let next_clause_start = if_stmt + .elif_else_clauses + .first() + .map(|clause| clause.range.start()); + if let Some(target) = build_unwrap_block_target( + if_stmt.range(), + if_stmt.range().start(), + &if_stmt.body, + next_clause_start, + source, + position, + ) { + return Some(target); + } + + for (index, clause) in if_stmt.elif_else_clauses.iter().enumerate() { + let next_clause_start = if_stmt + .elif_else_clauses + .get(index + 1) + .map(|next_clause| next_clause.range.start()); + if let Some(target) = build_unwrap_block_target( + if_stmt.range(), + clause.range.start(), + &clause.body, + next_clause_start, + source, + position, + ) { + return Some(target); + } + } + + None +} + +fn build_unwrap_block_target( + outer_range: TextRange, + header_start: TextSize, + body: &[Stmt], + next_clause_start: Option, + source: &str, + position: TextSize, +) -> Option { + let first_stmt = body.first()?; + let (to_indent, _) = line_indent_and_start(source, outer_range.start())?; + let (_, header_line_start) = line_indent_and_start(source, header_start)?; + let (from_indent, first_stmt_line_start) = + line_indent_and_start(source, first_stmt.range().start())?; + let colon_position = find_suite_colon(source, header_line_start, first_stmt_line_start)?; + if position < header_line_start || position > colon_position + TextSize::new(1) { + return None; + } + let block_start = line_end_position(source, colon_position); + if block_start > first_stmt.range().start() { + return None; + } + let block_end = if let Some(next_clause_start) = next_clause_start { + line_indent_and_start(source, next_clause_start)?.1 + } else { + line_end_position(source, outer_range.end()) + }; + if block_start >= block_end { + return None; + } + Some(UnwrapBlockTarget { + outer_range, + body_range: TextRange::new(block_start, block_end), + from_indent, + to_indent, + }) +} + +fn find_suite_colon( + source: &str, + header_line_start: TextSize, + body_line_start: TextSize, +) -> Option { + let header_start = header_line_start.to_usize(); + let header_end = body_line_start.to_usize().min(source.len()); + if header_start >= header_end { + return None; + } + let header = &source[header_start..header_end]; + let colon = find_suite_colon_in_header(header)?; + TextSize::try_from(header_start + colon).ok() +} + +fn find_suite_colon_in_header(header: &str) -> Option { + let mut quote = None::; + let mut escaped = false; + let mut in_comment = false; + let mut last_colon = None; + + for (index, ch) in header.char_indices() { + if in_comment { + if ch == '\n' { + in_comment = false; + } + continue; + } + if escaped { + escaped = false; + continue; + } + if let Some(state) = quote { + if state.triple { + if ch == state.delimiter + && header[index..].starts_with(triple_quote_delimiter(state.delimiter)) + { + quote = None; + } else if ch == '\\' { + escaped = true; + } + } else { + match ch { + '\\' => escaped = true, + current if current == state.delimiter => quote = None, + _ => {} + } + } + continue; + } + match ch { + '\'' | '"' => { + quote = Some(QuoteState { + delimiter: ch, + triple: header[index..].starts_with(triple_quote_delimiter(ch)), + }); + } + '#' => in_comment = true, + ':' => last_colon = Some(index), + _ => {} + } + } + + last_colon +} + +fn triple_quote_delimiter(ch: char) -> &'static str { + match ch { + '\'' => "'''", + '"' => "\"\"\"", + _ => unreachable!("only quote delimiters are supported"), + } +} + +fn find_try_clause_start( + source: &str, + body_start: TextSize, + outer_indent: &str, + keyword: &str, +) -> Option { + let mut line_end = body_start.to_usize().min(source.len()); + let expected_prefix = format!("{outer_indent}{keyword}"); + while line_end > 0 { + let search_end = line_end.saturating_sub(1); + let line_start = source[..search_end] + .rfind('\n') + .map(|index| index + 1) + .unwrap_or(0); + let line = source[line_start..line_end].trim_end_matches('\n'); + if line.starts_with(&expected_prefix) && find_suite_colon_in_header(line).is_some() { + return TextSize::try_from(line_start).ok(); + } + if line_start == 0 { + break; + } + line_end = line_start; + } + None +} + +fn range_contains_position(range: TextRange, position: TextSize) -> bool { + range.start() <= position && position < range.end() +} diff --git a/pyrefly/lib/test/lsp/code_actions.rs b/pyrefly/lib/test/lsp/code_actions.rs index 2b83c8698a..36ea5b5a57 100644 --- a/pyrefly/lib/test/lsp/code_actions.rs +++ b/pyrefly/lib/test/lsp/code_actions.rs @@ -367,6 +367,19 @@ fn apply_first_inline_parameter_action(code: &str) -> Option { Some(apply_refactor_edits_for_module(&module_info, &edits)) } +fn apply_first_unwrap_block_action(code: &str, selection: TextRange) -> Option { + 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 + .unwrap_block_code_actions(handle, selection) + .unwrap_or_default(); + let edits = actions.first()?.edits.clone(); + Some(apply_refactor_edits_for_module(&module_info, &edits)) +} + fn apply_first_safe_delete_action(code: &str) -> Option { apply_first_safe_delete_action_multi(&[("main", code)], "main") } @@ -431,6 +444,21 @@ fn assert_no_introduce_parameter_action(code: &str) { ); } +fn assert_no_unwrap_block_action(code: &str, selection: TextRange) { + 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 actions = transaction + .unwrap_block_code_actions(handle, selection) + .unwrap_or_default(); + assert!( + actions.is_empty(), + "expected no unwrap-block actions, found {}", + actions.len() + ); +} + fn compute_change_signature_actions( code: &str, ) -> ( @@ -5968,6 +5996,315 @@ class Child(Base): assert!(apply_first_safe_delete_action(code).is_none()); } +#[test] +fn unwrap_block_if_body() { + let code = r#" +def f(): + if True: + print(2) +"#; + let selection = find_nth_range(code, ":", 2); + let updated = + apply_first_unwrap_block_action(code, selection).expect("expected unwrap-block action"); + let expected = r#" +def f(): + print(2) +"#; + assert_eq!(expected.trim(), updated.trim()); +} + +#[test] +fn unwrap_block_if_header_selection() { + let code = r#" +def f(): + if True: + print(2) +"#; + let selection = find_nth_range(code, "True", 1); + let updated = + apply_first_unwrap_block_action(code, selection).expect("expected unwrap-block action"); + let expected = r#" +def f(): + print(2) +"#; + assert_eq!(expected.trim(), updated.trim()); +} + +#[test] +fn unwrap_block_if_else_branch_after_colon() { + let code = r#" +def f(): + if True: + print("then") + else: + # keep this comment + print("else") +"#; + let else_colon = find_nth_range(code, ":", 3); + let selection = TextRange::new(else_colon.end(), else_colon.end()); + let updated = + apply_first_unwrap_block_action(code, selection).expect("expected unwrap-block action"); + let expected = r#" +def f(): + # keep this comment + print("else") +"#; + assert_eq!(expected.trim(), updated.trim()); +} + +#[test] +fn unwrap_block_elif_branch() { + let code = r#" +def f(value): + if value == 0: + print("zero") + elif value == 1: + print("one") + else: + print("other") +"#; + let selection = find_nth_range(code, "value == 1", 1); + let updated = + apply_first_unwrap_block_action(code, selection).expect("expected unwrap-block action"); + let expected = r#" +def f(value): + print("one") +"#; + assert_eq!(expected.trim(), updated.trim()); +} + +#[test] +fn unwrap_block_rejects_loop_else() { + let code = r#" +def f(items): + for item in items: + print(item) + else: + print("done") +"#; + let selection = find_nth_range(code, ":", 3); + assert_no_unwrap_block_action(code, selection); +} + +#[test] +fn unwrap_block_with_statement() { + let code = r#" +def f(): + with open("x.txt") as f: + print(f.read()) +"#; + let selection = find_nth_range(code, ":", 2); + let updated = + apply_first_unwrap_block_action(code, selection).expect("expected unwrap-block action"); + let expected = r#" +def f(): + print(f.read()) +"#; + assert_eq!(expected.trim(), updated.trim()); +} + +#[test] +fn unwrap_block_while_without_else() { + let code = r#" +def f(): + while True: + print("tick") +"#; + let selection = find_nth_range(code, ":", 2); + let updated = + apply_first_unwrap_block_action(code, selection).expect("expected unwrap-block action"); + let expected = r#" +def f(): + print("tick") +"#; + assert_eq!(expected.trim(), updated.trim()); +} + +#[test] +fn unwrap_block_outer_nested_if() { + let code = r#" +def f(): + if True: + if True: + print("value") +"#; + let selection = find_nth_range(code, ":", 2); + let updated = + apply_first_unwrap_block_action(code, selection).expect("expected unwrap-block action"); + let expected = r#" +def f(): + if True: + print("value") +"#; + assert_eq!(expected.trim(), updated.trim()); +} + +#[test] +fn unwrap_block_inner_nested_if() { + let code = r#" +def f(): + if True: + if True: + print("value") +"#; + let selection = find_nth_range(code, ":", 3); + let updated = + apply_first_unwrap_block_action(code, selection).expect("expected unwrap-block action"); + let expected = r#" +def f(): + if True: + print("value") +"#; + assert_eq!(expected.trim(), updated.trim()); +} + +#[test] +fn unwrap_block_multiline_header() { + let code = r#" +def f(): + if ( + True + ): + print("then") +"#; + let selection = find_nth_range(code, ":", 2); + let updated = + apply_first_unwrap_block_action(code, selection).expect("expected unwrap-block action"); + let expected = r#" +def f(): + print("then") +"#; + assert_eq!(expected.trim(), updated.trim()); +} + +#[test] +fn unwrap_block_multiline_triple_quoted_string_header() { + let code = r#" +def f(): + if bool("""first: +second"""): + print("then") +"#; + let selection = find_nth_range(code, "second", 1); + let updated = + apply_first_unwrap_block_action(code, selection).expect("expected unwrap-block action"); + let expected = r#" +def f(): + print("then") +"#; + assert_eq!(expected.trim(), updated.trim()); +} + +#[test] +fn unwrap_block_header_comment_with_colon() { + let code = r#" +def f(): + if True: # comment with : + print("then") +"#; + let selection = find_nth_range(code, ":", 2); + let updated = + apply_first_unwrap_block_action(code, selection).expect("expected unwrap-block action"); + let expected = r#" +def f(): + print("then") +"#; + assert_eq!(expected.trim(), updated.trim()); +} + +#[test] +fn unwrap_block_rejects_comment_colon_selection() { + let code = r#" +def f(): + if True: # comment with : + print("then") +"#; + let comment_colon = + TextSize::try_from(code.rfind(':').expect("expected comment colon")).unwrap(); + let selection = TextRange::new(comment_colon, comment_colon); + assert_no_unwrap_block_action(code, selection); +} + +#[test] +fn unwrap_block_f_string_format_spec() { + let code = r#" +def f(value): + if f"{value:02}": + print("then") +"#; + let selection = find_nth_range(code, "value:02", 1); + let updated = + apply_first_unwrap_block_action(code, selection).expect("expected unwrap-block action"); + let expected = r#" +def f(value): + print("then") +"#; + assert_eq!(expected.trim(), updated.trim()); +} + +#[test] +fn unwrap_block_try_selected_except_branch() { + let code = r#" +def f(value): + try: + print("try") + except ValueError: + print("value") + except TypeError: + print("type") +"#; + let selection = find_nth_range(code, ":", 4); + let updated = + apply_first_unwrap_block_action(code, selection).expect("expected unwrap-block action"); + let expected = r#" +def f(value): + print("type") +"#; + assert_eq!(expected.trim(), updated.trim()); +} + +#[test] +fn unwrap_block_try_finally_branch() { + let code = r#" +def f(): + try: + print("try") + finally: + print("cleanup") +"#; + let selection = find_nth_range(code, ":", 3); + let updated = + apply_first_unwrap_block_action(code, selection).expect("expected unwrap-block action"); + let expected = r#" +def f(): + print("cleanup") +"#; + assert_eq!(expected.trim(), updated.trim()); +} + +#[test] +fn unwrap_block_rejects_try_header_with_finally() { + let code = r#" +def f(): + try: + print("try") + finally: + print("cleanup") +"#; + let selection = find_nth_range(code, "try", 1); + assert_no_unwrap_block_action(code, selection); +} + +#[test] +fn unwrap_block_rejects_single_line_suite() { + let code = r#" +def f(): + if True: print("then") +"#; + let selection = find_nth_range(code, ":", 2); + assert_no_unwrap_block_action(code, selection); +} + #[test] fn pytest_fixture_type_annotation_code_actions() { let conftest = r#"