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::CompletionItemTextEdit;
 use lsp_types::InsertTextFormat;
 use lsp_types::TextEdit;
 use pyrefly_build::handle::Handle;
@@ -354,6 +355,150 @@ impl Transaction<'_> {
         }
     }
 
+    /// Offers to insert the closer after an opening triple-quoted string.
+    ///
+    /// Only a run of exactly three quotes is treated as that opener. A single or
+    /// double quote is already an open string (Literal values, dict keys, kwargs),
+    /// so expanding those would steal the completions that belong in that context.
+    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 != 3 {
+            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;
+                }
+            }
+        }
+        if Self::is_inside_string_or_comment(&before[..quote_start]) {
+            return;
+        }
+        let closer = quote.to_string().repeat(3);
+        if source[pos..].starts_with(&closer) {
+            return;
+        }
+        let new_text = if supports_snippets {
+            format!("$0{closer}")
+        } else {
+            closer.clone()
+        };
+        let range = module_info.to_lsp_range(TextRange::new(position, 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(CompletionItemTextEdit::TextEdit(TextEdit {
+                range,
+                new_text,
+            })),
+            ..Default::default()
+        }));
+    }
+
+    /// Returns true if the end of `source` lies inside a string literal or a comment,
+    /// so quotes typed there close a string or are comment text rather than open one.
+    fn is_inside_string_or_comment(source: &str) -> bool {
+        let bytes = source.as_bytes();
+        let mut i = 0;
+        let mut in_comment = false;
+        let mut string: Option<(u8, bool)> = None;
+        while i < bytes.len() {
+            let b = bytes[i];
+            if in_comment {
+                if b == b'\n' {
+                    in_comment = false;
+                }
+                i += 1;
+                continue;
+            }
+            match string {
+                Some((q, triple)) => {
+                    if b == b'\\' {
+                        i += 2;
+                        continue;
+                    }
+                    if triple {
+                        if bytes[i..].starts_with(&[q, q, q]) {
+                            string = None;
+                            i += 3;
+                            continue;
+                        }
+                    } else if b == q || b == b'\n' {
+                        string = None;
+                    }
+                    i += 1;
+                }
+                None => {
+                    if b == b'#' {
+                        in_comment = true;
+                        i += 1;
+                    } else if b == b'"' || b == b'\'' {
+                        if bytes[i..].starts_with(&[b, b, b]) {
+                            string = Some((b, true));
+                            i += 3;
+                        } else {
+                            string = Some((b, false));
+                            i += 1;
+                        }
+                    } else {
+                        i += 1;
+                    }
+                }
+            }
+        }
+        in_comment || string.is_some()
+    }
+
     /// Retrieves documentation for an export to display in completion items.
     fn get_documentation_from_export(
         &self,
@@ -1097,6 +1242,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);