diff --git a/pyrefly/lib/lsp/wasm/completion.rs b/pyrefly/lib/lsp/wasm/completion.rs index ec4b839985..cf6bd2b7ce 100644 --- a/pyrefly/lib/lsp/wasm/completion.rs +++ b/pyrefly/lib/lsp/wasm/completion.rs @@ -12,6 +12,7 @@ use lsp_types::CompletionItem; use lsp_types::CompletionItemKind; use lsp_types::CompletionItemLabelDetails; use lsp_types::CompletionItemTag; +use lsp_types::CompletionTextEdit; use lsp_types::InsertTextFormat; use lsp_types::TextEdit; use pyrefly_build::handle::Handle; @@ -350,6 +351,96 @@ impl Transaction<'_> { } } + /// Offers to close (or expand into) a triple-quoted string after the opening quotes. + fn add_triple_quoted_string_completions( + &self, + handle: &Handle, + position: TextSize, + supports_snippets: bool, + completions: &mut Vec, + ) { + let Some(module_info) = self.get_module_info(handle) else { + return; + }; + let source = module_info.contents(); + let pos = position.to_usize(); + if pos == 0 || pos > source.len() { + return; + } + let before = &source[..pos]; + let Some(quote) = before.chars().next_back() else { + return; + }; + if quote != '"' && quote != '\'' { + return; + } + let mut count = 0usize; + for c in before.chars().rev() { + if c == quote && count < 3 { + count += 1; + } else { + break; + } + } + if count == 0 { + return; + } + let quote_start = pos - count; + if quote_start > 0 && before.as_bytes()[quote_start - 1] == quote as u8 { + return; + } + if quote_start > 0 { + let prev = before.as_bytes()[quote_start - 1]; + if prev.is_ascii_alphanumeric() || prev == b'_' { + let mut ident_start = quote_start; + while ident_start > 0 { + let b = before.as_bytes()[ident_start - 1]; + if b.is_ascii_alphanumeric() || b == b'_' { + ident_start -= 1; + } else { + break; + } + } + let ident = before[ident_start..quote_start].to_ascii_lowercase(); + let is_string_prefix = matches!( + ident.as_str(), + "r" | "u" | "f" | "b" | "fr" | "rf" | "br" | "rb" + ); + if !is_string_prefix { + return; + } + } + } + let closer = quote.to_string().repeat(3); + if source[pos..].starts_with(&closer) { + return; + } + let (edit_start, new_text) = if count < 3 { + let opener_and_closer = if supports_snippets { + format!("{closer}$0{closer}") + } else { + format!("{closer}{closer}") + }; + (quote_start, opener_and_closer) + } else if supports_snippets { + (pos, format!("$0{closer}")) + } else { + (pos, closer.clone()) + }; + let Some(start) = TextSize::try_from(edit_start).ok() else { + return; + }; + let range = module_info.to_lsp_range(TextRange::new(start, position)); + completions.push(RankedCompletion::new(CompletionItem { + label: closer, + detail: Some("triple-quoted string".to_owned()), + kind: Some(CompletionItemKind::SNIPPET), + insert_text_format: supports_snippets.then_some(InsertTextFormat::SNIPPET), + text_edit: Some(CompletionTextEdit::Edit(TextEdit { range, new_text })), + ..Default::default() + })); + } + /// Retrieves documentation for an export to display in completion items. fn get_documentation_from_export( &self, @@ -1056,6 +1147,12 @@ impl Transaction<'_> { auto_import, } = options; let mut result: Vec = Vec::new(); + self.add_triple_quoted_string_completions( + handle, + position, + supports_snippet_completions, + &mut result, + ); let mut is_incomplete = false; let mut allow_function_call_parens = false; let ast = self.get_ast(handle); diff --git a/pyrefly/lib/test/lsp/completion.rs b/pyrefly/lib/test/lsp/completion.rs index 399ab23286..e05794d3df 100644 --- a/pyrefly/lib/test/lsp/completion.rs +++ b/pyrefly/lib/test/lsp/completion.rs @@ -13,6 +13,7 @@ use pyrefly_build::handle::Handle; use pyrefly_python::sys_info::PythonVersion; use ruff_text_size::TextSize; +use crate::lsp::wasm::completion::CompletionOptions; use crate::state::lsp::ImportFormat; use crate::state::require::Require; use crate::state::state::State; @@ -4195,3 +4196,131 @@ def users() -> None: "missing status_code in completions:\n{trimmed}" ); } + +fn get_triple_quoted_string_report( + supports_snippets: bool, +) -> impl Fn(&State, &Handle, TextSize) -> String { + move |state: &State, handle: &Handle, position: TextSize| { + let mut report = "Completion Results:".to_owned(); + for item in state + .transaction() + .completion_with_incomplete( + handle, + position, + ImportFormat::Absolute, + CompletionOptions { + supports_snippet_completions: supports_snippets, + ..Default::default() + }, + None, + ) + .0 + { + if item.detail.as_deref() != Some("triple-quoted string") { + continue; + } + report.push_str("\n- ("); + report.push_str(&format!("{:?}", item.kind.unwrap())); + report.push_str(") "); + report.push_str(&item.label); + if let Some(detail) = item.detail { + report.push_str(": "); + report.push_str(&detail); + } + if let Some(text_edit) = item.text_edit { + report.push_str(" with text edit: "); + report.push_str(&format!("{:?}", text_edit)); + } + if let Some(fmt) = item.insert_text_format { + report.push_str(&format!(" format={fmt:?}")); + } + } + report + } +} + +#[test] +fn completion_closes_triple_quoted_string() { + let code = r#" +x = """ +# ^ +"#; + let report = get_batched_lsp_operations_report_allow_error( + &[("main", code)], + get_triple_quoted_string_report(true), + ); + let trimmed = report.trim(); + assert!( + trimmed.contains("triple-quoted string"), + "missing triple-quoted string completion:\n{trimmed}" + ); + assert!( + trimmed.contains("new_text: \"$0"), + "expected snippet closer after opening quotes:\n{trimmed}" + ); +} + +#[test] +fn completion_expands_partial_quotes_to_triple_quoted_string() { + let code = r#" +x = " +# ^ +"#; + let report = get_batched_lsp_operations_report_allow_error( + &[("main", code)], + get_triple_quoted_string_report(true), + ); + let trimmed = report.trim(); + assert!( + trimmed.contains("new_text: \"\\\"\\\"\\\"$0"), + "expected expansion to a triple-quoted pair:\n{trimmed}" + ); +} + +#[test] +fn completion_skips_triple_quoted_string_when_closer_exists() { + let code = r#" +x = """""" +# ^ +"#; + let report = get_batched_lsp_operations_report_allow_error( + &[("main", code)], + get_triple_quoted_string_report(false), + ); + assert!( + !report.contains("triple-quoted string"), + "should not offer a closer when one is already present:\n{report}" + ); +} + +#[test] +fn completion_skips_quotes_after_identifier() { + let code = r#" +foo" +# ^ +"#; + let report = get_batched_lsp_operations_report_allow_error( + &[("main", code)], + get_triple_quoted_string_report(false), + ); + assert!( + !report.contains("triple-quoted string"), + "should not treat identifier-adjacent quotes as a string opener:\n{report}" + ); +} + +#[test] +fn completion_closes_prefixed_triple_quoted_string() { + let code = r#" +x = f""" +# ^ +"#; + let report = get_batched_lsp_operations_report_allow_error( + &[("main", code)], + get_triple_quoted_string_report(false), + ); + assert!( + report.contains("triple-quoted string"), + "missing closer after an f-string prefix:\n{report}" + ); +}