Offer TypedDict keys as completions after a dot

pyrefly/lib/lsp/wasm/completion.rs

@@ -1362,7 +1362,7 @@ impl Transaction<'_> {
                 }
             }
             Some(IdentifierWithContext {
-                identifier: _,
+                identifier,
                 context: IdentifierContext::Attribute { base_range, .. },
             }) => {
                 let expected_type = covering_nodes
@@ -1396,6 +1396,21 @@ impl Transaction<'_> {
                             }
                         }
                     }
+                    // Value patterns only allow dotted names, so a subscript would be invalid.
+                    if let Some(nodes) = covering_nodes.as_deref()
+                        && !nodes
+                            .iter()
+                            .any(|node| matches!(node, AnyNodeRef::PatternMatchValue(_)))
+                    {
+                        self.add_typed_dict_attribute_key_completions(
+                            handle,
+                            &base_type,
+                            base_range,
+                            &identifier,
+                            nodes,
+                            &mut result,
+                        );
+                    }
                     self.add_attribute_completions_for_type(
                         handle,
                         base_type,

pyrefly/lib/state/lsp/dict_completions.rs

@@ -10,6 +10,8 @@ use std::collections::BTreeSet;
 
 use lsp_types::CompletionItem;
 use lsp_types::CompletionItemKind;
+use lsp_types::CompletionItemTextEdit;
+use lsp_types::TextEdit;
 use pyrefly_build::handle::Handle;
 use pyrefly_python::ast::Ast;
 use pyrefly_python::short_identifier::ShortIdentifier;
@@ -25,9 +27,12 @@ use ruff_python_ast::ExprStringLiteral;
 use ruff_python_ast::Identifier;
 use ruff_python_ast::Keyword;
 use ruff_python_ast::ModModule;
+use ruff_python_ast::StringFlags;
+use ruff_python_ast::str::Quote;
 use ruff_text_size::Ranged;
 use ruff_text_size::TextRange;
 use ruff_text_size::TextSize;
+use serde_json::to_string;
 
 use crate::alt::answers_solver::AnswersSolver;
 use crate::alt::polars_specials::polars_function_treats_strings_as_columns;
@@ -1095,6 +1100,78 @@ impl<'a> Transaction<'a> {
         }
     }
 
+    /// Offers the keys of a TypedDict after `base.` as completions that rewrite the
+    /// attribute access `base.key` into the subscript `base["key"]`.
+    pub(crate) fn add_typed_dict_attribute_key_completions(
+        &self,
+        handle: &Handle,
+        base_type: &Type,
+        base_range: TextRange,
+        attribute: &Identifier,
+        covering_nodes: &[AnyNodeRef],
+        completions: &mut Vec<RankedCompletion>,
+    ) {
+        if !Self::type_contains_typed_dict(base_type) {
+            return;
+        }
+        let Some(module_info) = self.get_module_info(handle) else {
+            return;
+        };
+        // The base range excludes grouping parentheses, so `(a).f` has a `)` before the dot.
+        let before_attribute = TextRange::new(base_range.end(), attribute.start());
+        let dot_offset = module_info
+            .code_at(before_attribute)
+            .rfind('.')
+            .expect("An attribute access has a dot before its name");
+        let dot = before_attribute.start() + TextSize::new(dot_offset as u32);
+        let edit_range = module_info.to_lsp_range(TextRange::new(dot, attribute.end()));
+        // LSP requires a completion's text edit to be on a single line.
+        if edit_range.start.line != edit_range.end.line {
+            return;
+        }
+        // Before Python 3.12 (PEP 701), an f-string replacement field cannot contain a
+        // backslash or the quote of any enclosing f-string.
+        let f_string_quotes: Vec<Quote> = covering_nodes
+            .iter()
+            .filter_map(|node| match node {
+                AnyNodeRef::FString(f_string) => Some(f_string.flags.quote_style()),
+                _ => None,
+            })
+            .collect();
+        let Some(quote) = [Quote::Double, Quote::Single]
+            .into_iter()
+            .find(|quote| !f_string_quotes.contains(quote))
+            .map(Quote::as_char)
+        else {
+            return;
+        };
+        let Some(keys) = self.collect_typed_dict_keys(handle, base_type.clone()) else {
+            return;
+        };
+        // Clients filter on the text from the start of the edit, which includes the dot.
+        let accessor = module_info.code_at(TextRange::new(dot, attribute.start()));
+        for (key, ty) in keys {
+            let quoted_key = if f_string_quotes.is_empty() {
+                to_string(&key).expect("Serializing a string cannot fail")
+            } else if key.contains(['"', '\'', '\\']) || key.chars().any(char::is_control) {
+                continue;
+            } else {
+                format!("{quote}{key}{quote}")
+            };
+            completions.push(RankedCompletion::new(CompletionItem {
+                filter_text: Some(format!("{accessor}{key}")),
+                label: key,
+                detail: Some(ty.to_string()),
+                kind: Some(CompletionItemKind::Field),
+                text_edit: Some(CompletionItemTextEdit::TextEdit(TextEdit {
+                    range: edit_range,
+                    new_text: format!("[{quoted_key}]"),
+                })),
+                ..Default::default()
+            }));
+        }
+    }
+
     /// Returns unprovided contextual fields at a builtin dict keyword-name position.
     pub(crate) fn typed_dict_constructor_kwargs(
         &self,