feat(lsp): complete closing triple-quoted strings

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<RankedCompletion>,
+    ) {
+        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<RankedCompletion> = 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);

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}"
+    );
+}