From 2a73b62215f0243a9567cd37eadf72545544461f Mon Sep 17 00:00:00 2001 From: Asuka Minato Date: Mon, 13 Apr 2026 02:42:42 +0900 Subject: [PATCH 1/5] fix --- pyrefly/lib/lsp/non_wasm/server.rs | 4 + pyrefly/lib/state/lsp.rs | 8 + pyrefly/lib/state/lsp/quick_fixes/mod.rs | 1 + .../lib/state/lsp/quick_fixes/unwrap_block.rs | 295 ++++++++++++++++++ pyrefly/lib/test/lsp/code_actions.rs | 80 +++++ 5 files changed, 388 insertions(+) create mode 100644 pyrefly/lib/state/lsp/quick_fixes/unwrap_block.rs diff --git a/pyrefly/lib/lsp/non_wasm/server.rs b/pyrefly/lib/lsp/non_wasm/server.rs index 9cdf981f61..cd38845510 100644 --- a/pyrefly/lib/lsp/non_wasm/server.rs +++ b/pyrefly/lib/lsp/non_wasm/server.rs @@ -4585,6 +4585,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 1caeec82c0..81d63c2aae 100644 --- a/pyrefly/lib/state/lsp.rs +++ b/pyrefly/lib/state/lsp.rs @@ -2994,6 +2994,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) + } + /// Determines whether a module is a third-party package. /// /// Checks if the module's path is located within any of the configured diff --git a/pyrefly/lib/state/lsp/quick_fixes/mod.rs b/pyrefly/lib/state/lsp/quick_fixes/mod.rs index 16331113d4..a4a67ccdf0 100644 --- a/pyrefly/lib/state/lsp/quick_fixes/mod.rs +++ b/pyrefly/lib/state/lsp/quick_fixes/mod.rs @@ -26,3 +26,4 @@ pub(crate) mod pytest_fixture; pub(crate) mod redundant_cast; pub(crate) mod safe_delete; pub(crate) mod types; +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..b17442e1e4 --- /dev/null +++ b/pyrefly/lib/state/lsp/quick_fixes/unwrap_block.rs @@ -0,0 +1,295 @@ +/* + * 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, +} + +/// 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::REFACTOR_REWRITE, + }]) +} + +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) => { + 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); + } + find_unwrap_block_target_in_body(&try_stmt.finalbody, 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 (from_indent, first_stmt_line_start) = + line_indent_and_start(source, first_stmt.range().start())?; + let colon_position = find_suite_colon(source, header_start, first_stmt_line_start, &to_indent)?; + if position != colon_position && 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_start: TextSize, + body_line_start: TextSize, + outer_indent: &str, +) -> Option { + let mut line_end = body_line_start.to_usize().min(source.len()); + let min_start = header_start.to_usize(); + while line_end > min_start { + 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.trim().is_empty() { + let indent_len = line + .chars() + .take_while(|ch| *ch == ' ' || *ch == '\t') + .map(char::len_utf8) + .sum::(); + if line.get(..indent_len) == Some(outer_indent) && line.rfind(':').is_some() { + let colon = line.rfind(':').expect("checked above"); + return TextSize::try_from(line_start + colon).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 77b46086cb..d0593db234 100644 --- a/pyrefly/lib/test/lsp/code_actions.rs +++ b/pyrefly/lib/test/lsp/code_actions.rs @@ -342,6 +342,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") } @@ -406,6 +419,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 assert_no_extract_variable_action(code: &str) { let (_, actions, _) = compute_extract_variable_actions(code); assert!( @@ -4574,6 +4602,58 @@ 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_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_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 pytest_fixture_type_annotation_code_actions() { let conftest = r#" From 34d7c19652764aa5dbf468be71557507eb055b82 Mon Sep 17 00:00:00 2001 From: Asuka Minato Date: Sat, 18 Apr 2026 13:30:56 +0900 Subject: [PATCH 2/5] Update pyrefly/lib/test/lsp/code_actions.rs Co-authored-by: Copilot <175728472+Copilot@users.noreply.github.com> --- pyrefly/lib/test/lsp/code_actions.rs | 70 ++++++++++++++++++++++++++++ 1 file changed, 70 insertions(+) diff --git a/pyrefly/lib/test/lsp/code_actions.rs b/pyrefly/lib/test/lsp/code_actions.rs index d0593db234..3f8196133b 100644 --- a/pyrefly/lib/test/lsp/code_actions.rs +++ b/pyrefly/lib/test/lsp/code_actions.rs @@ -4654,6 +4654,76 @@ def f(items): 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_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_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 pytest_fixture_type_annotation_code_actions() { let conftest = r#" From 2887be89a4d97bdcd1acfec6a3af5bb28cc0c553 Mon Sep 17 00:00:00 2001 From: Asuka Minato Date: Sat, 18 Apr 2026 13:39:57 +0900 Subject: [PATCH 3/5] comment --- .../lib/state/lsp/quick_fixes/unwrap_block.rs | 111 ++++++++++++++++-- pyrefly/lib/test/lsp/code_actions.rs | 63 ++++++++++ 2 files changed, 163 insertions(+), 11 deletions(-) diff --git a/pyrefly/lib/state/lsp/quick_fixes/unwrap_block.rs b/pyrefly/lib/state/lsp/quick_fixes/unwrap_block.rs index b17442e1e4..b5dbbb19fb 100644 --- a/pyrefly/lib/state/lsp/quick_fixes/unwrap_block.rs +++ b/pyrefly/lib/state/lsp/quick_fixes/unwrap_block.rs @@ -156,7 +156,72 @@ fn find_unwrap_block_target_in_stmt( { return Some(target); } - find_unwrap_block_target_in_body(&try_stmt.finalbody, source, position) + 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() + .map(Ranged::range) + .map(|range| range.start()); + if let Some(target) = build_unwrap_block_target( + stmt.range(), + stmt.range().start(), + &try_stmt.orelse, + try_stmt + .finalbody + .first() + .map(Ranged::range) + .map(|range| range.start()), + source, + position, + ) { + return Some(target); + } + build_unwrap_block_target( + stmt.range(), + else_start.unwrap_or(stmt.range().start()), + &try_stmt.finalbody, + None, + source, + position, + ) } Stmt::Match(match_stmt) => { for case in &match_stmt.cases { @@ -271,16 +336,11 @@ fn find_suite_colon( .map(|index| index + 1) .unwrap_or(0); let line = source[line_start..line_end].trim_end_matches('\n'); - if !line.trim().is_empty() { - let indent_len = line - .chars() - .take_while(|ch| *ch == ' ' || *ch == '\t') - .map(char::len_utf8) - .sum::(); - if line.get(..indent_len) == Some(outer_indent) && line.rfind(':').is_some() { - let colon = line.rfind(':').expect("checked above"); - return TextSize::try_from(line_start + colon).ok(); - } + if !line.trim().is_empty() + && line.starts_with(outer_indent) + && let Some(colon) = find_suite_colon_in_line(line) + { + return TextSize::try_from(line_start + colon).ok(); } if line_start == 0 { break; @@ -290,6 +350,35 @@ fn find_suite_colon( None } +fn find_suite_colon_in_line(line: &str) -> Option { + let mut quote = None; + let mut escaped = false; + let mut last_colon = None; + + for (index, ch) in line.char_indices() { + if escaped { + escaped = false; + continue; + } + if let Some(quote_char) = quote { + match ch { + '\\' => escaped = true, + current if current == quote_char => quote = None, + _ => {} + } + continue; + } + match ch { + '\'' | '"' => quote = Some(ch), + '#' => break, + ':' => last_colon = Some(index), + _ => {} + } + } + + last_colon +} + 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 3f8196133b..af0980aa3a 100644 --- a/pyrefly/lib/test/lsp/code_actions.rs +++ b/pyrefly/lib/test/lsp/code_actions.rs @@ -4724,6 +4724,69 @@ def f(): 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_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_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#" From 0b82129c221a75c730a2fba93caab5ab93808658 Mon Sep 17 00:00:00 2001 From: Asuka Minato Date: Tue, 21 Apr 2026 13:30:21 +0900 Subject: [PATCH 4/5] comment --- .../lib/state/lsp/quick_fixes/unwrap_block.rs | 155 ++++++++++++------ pyrefly/lib/test/lsp/code_actions.rs | 124 ++++++++++++++ 2 files changed, 227 insertions(+), 52 deletions(-) diff --git a/pyrefly/lib/state/lsp/quick_fixes/unwrap_block.rs b/pyrefly/lib/state/lsp/quick_fixes/unwrap_block.rs index b5dbbb19fb..1dbd980a94 100644 --- a/pyrefly/lib/state/lsp/quick_fixes/unwrap_block.rs +++ b/pyrefly/lib/state/lsp/quick_fixes/unwrap_block.rs @@ -30,6 +30,12 @@ struct UnwrapBlockTarget { 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<'_>, @@ -139,6 +145,7 @@ fn find_unwrap_block_target_in_stmt( ) } 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); @@ -195,28 +202,29 @@ fn find_unwrap_block_target_in_stmt( return Some(target); } } - let else_start = try_stmt - .orelse - .first() - .map(Ranged::range) - .map(|range| range.start()); - if let Some(target) = build_unwrap_block_target( - stmt.range(), - stmt.range().start(), - &try_stmt.orelse, - try_stmt - .finalbody - .first() - .map(Ranged::range) - .map(|range| range.start()), - source, - position, - ) { + 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(), - else_start.unwrap_or(stmt.range().start()), + finally_start?, &try_stmt.finalbody, None, source, @@ -295,10 +303,11 @@ fn build_unwrap_block_target( ) -> 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_start, first_stmt_line_start, &to_indent)?; - if position != colon_position && position != colon_position + TextSize::new(1) { + 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); @@ -323,54 +332,62 @@ fn build_unwrap_block_target( fn find_suite_colon( source: &str, - header_start: TextSize, + header_line_start: TextSize, body_line_start: TextSize, - outer_indent: &str, ) -> Option { - let mut line_end = body_line_start.to_usize().min(source.len()); - let min_start = header_start.to_usize(); - while line_end > min_start { - 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.trim().is_empty() - && line.starts_with(outer_indent) - && let Some(colon) = find_suite_colon_in_line(line) - { - return TextSize::try_from(line_start + colon).ok(); - } - if line_start == 0 { - break; - } - line_end = line_start; + 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; } - 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_line(line: &str) -> Option { - let mut quote = None; +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 line.char_indices() { + 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(quote_char) = quote { - match ch { - '\\' => escaped = true, - current if current == quote_char => quote = None, - _ => {} + 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(ch), - '#' => break, + '\'' | '"' => { + quote = Some(QuoteState { + delimiter: ch, + triple: header[index..].starts_with(triple_quote_delimiter(ch)), + }); + } + '#' => in_comment = true, ':' => last_colon = Some(index), _ => {} } @@ -379,6 +396,40 @@ fn find_suite_colon_in_line(line: &str) -> Option { 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 af0980aa3a..1c5d636930 100644 --- a/pyrefly/lib/test/lsp/code_actions.rs +++ b/pyrefly/lib/test/lsp/code_actions.rs @@ -4619,6 +4619,23 @@ def f(): 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#" @@ -4641,6 +4658,27 @@ def f(): 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#" @@ -4688,6 +4726,44 @@ def f(): 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#" @@ -4707,6 +4783,24 @@ def f(): 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#" @@ -4737,6 +4831,23 @@ def f(): 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#" @@ -4777,6 +4888,19 @@ def f(): 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#" From d5ea5b631cf11b870d974c87f5da0a19ec041454 Mon Sep 17 00:00:00 2001 From: Asuka Minato Date: Sat, 31 Oct 2026 03:00:07 +0900 Subject: [PATCH 5/5] fix api --- pyrefly/lib/state/lsp/quick_fixes/unwrap_block.rs | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/pyrefly/lib/state/lsp/quick_fixes/unwrap_block.rs b/pyrefly/lib/state/lsp/quick_fixes/unwrap_block.rs index 1dbd980a94..11dcdca58c 100644 --- a/pyrefly/lib/state/lsp/quick_fixes/unwrap_block.rs +++ b/pyrefly/lib/state/lsp/quick_fixes/unwrap_block.rs @@ -56,7 +56,7 @@ pub(crate) fn unwrap_block_code_actions( Some(vec![LocalRefactorCodeAction { title: "Unwrap block".to_owned(), edits: vec![(module_info.dupe(), removal_range, replacement)], - kind: CodeActionKind::REFACTOR_REWRITE, + kind: CodeActionKind::RefactorRewrite, }]) }