From fc27ee2e2baf59d7d35eeb2d9916d85809b01822 Mon Sep 17 00:00:00 2001 From: Parman Mohammadalizadeh Date: Sun, 1 Nov 2026 11:37:16 +0100 Subject: [PATCH 1/4] Add failing tests for TypedDict key completion after a dot Typing `a.` on a TypedDict instance offers only the mapping methods, so the keys can only be discovered by switching to `a["`. These tests pin the behavior requested in #5478. Each key is offered after the dot, and accepting it rewrites the access into `a["key"]`, replacing the dot and any typed prefix. Inherited and NotRequired/ReadOnly keys are included, keys that are not identifiers are quoted correctly, dict methods still appear, and plain classes and the TypedDict class object get no rewriting items. --- pyrefly/lib/test/lsp/completion.rs | 170 +++++++++++++++++++++++++++++ 1 file changed, 170 insertions(+) diff --git a/pyrefly/lib/test/lsp/completion.rs b/pyrefly/lib/test/lsp/completion.rs index 5725ebc119..9b7862eda6 100644 --- a/pyrefly/lib/test/lsp/completion.rs +++ b/pyrefly/lib/test/lsp/completion.rs @@ -8,6 +8,10 @@ use lsp_types::CompletionItem; use lsp_types::CompletionItemKind; use lsp_types::CompletionItemTag; +use lsp_types::CompletionItemTextEdit; +use lsp_types::Position; +use lsp_types::Range; +use lsp_types::TextEdit; use pretty_assertions::assert_eq; use pyrefly_build::handle::Handle; use pyrefly_python::sys_info::PythonVersion; @@ -2917,6 +2921,172 @@ a.zz ); } +fn completions_at_cursor(code: &str) -> Vec { + let (handles, state) = mk_multi_file_state(&[("main", code)], Require::Exports, false); + let position = extract_cursors_for_test(code)[0]; + state.transaction().completion( + &handles["main"], + position, + ImportFormat::Absolute, + true, + None, + ) +} + +/// The labels of the completions that rewrite the attribute access into a subscript. +fn subscript_rewrite_labels(code: &str) -> Vec { + completions_at_cursor(code) + .into_iter() + .filter(|item| item.text_edit.is_some()) + .map(|item| item.label) + .collect() +} + +// https://github.com/meta/pyrefly/issues/5478 +#[test] +fn typed_dict_attribute_completion_rewrites_to_subscript() { + let code = r#" +from typing import TypedDict + +class A(TypedDict): + f1: int + f2: str + +a: A +a. +# ^ +"#; + let completions = completions_at_cursor(code); + let f1 = completions + .iter() + .find(|item| item.label == "f1") + .expect("missing completion for TypedDict key `f1`"); + assert_eq!(f1.kind, Some(CompletionItemKind::Field)); + assert_eq!(f1.detail.as_deref(), Some("int")); + assert_eq!(f1.filter_text.as_deref(), Some(".f1")); + assert_eq!( + f1.text_edit, + Some(CompletionItemTextEdit::TextEdit(TextEdit { + range: Range::new(Position::new(8, 1), Position::new(8, 2)), + new_text: r#"["f1"]"#.to_owned(), + })) + ); + assert_eq!(subscript_rewrite_labels(code), vec!["f1", "f2"]); + for method in ["get", "keys", "items"] { + assert!( + completions + .iter() + .any(|item| item.label == method && item.kind == Some(CompletionItemKind::Method)), + "missing dict method `{method}`" + ); + } +} + +#[test] +fn typed_dict_attribute_completion_replaces_typed_prefix() { + let code = r#" +from typing import TypedDict + +class A(TypedDict): + field: int + +a: A +a.fi +# ^ +"#; + let completions = completions_at_cursor(code); + let field = completions + .iter() + .find(|item| item.label == "field") + .expect("missing completion for TypedDict key `field`"); + assert_eq!(field.filter_text.as_deref(), Some(".field")); + assert_eq!( + field.text_edit, + Some(CompletionItemTextEdit::TextEdit(TextEdit { + range: Range::new(Position::new(7, 1), Position::new(7, 4)), + new_text: r#"["field"]"#.to_owned(), + })) + ); +} + +#[test] +fn typed_dict_attribute_completion_includes_inherited_and_qualified_keys() { + let code = r#" +from typing_extensions import NotRequired, ReadOnly, TypedDict + +class Base(TypedDict): + base: int + +class Child(Base): + child: str + optional: NotRequired[int] + frozen: ReadOnly[bytes] + +c: Child +c. +# ^ +"#; + assert_eq!( + subscript_rewrite_labels(code), + vec!["base", "child", "frozen", "optional"] + ); +} + +#[test] +fn typed_dict_attribute_completion_quotes_non_identifier_keys() { + let code = r#" +from typing import TypedDict + +Weird = TypedDict("Weird", {"two words": int, 'say "hi"': str}) + +w: Weird +w. +# ^ +"#; + let new_texts: Vec<_> = completions_at_cursor(code) + .into_iter() + .filter_map(|item| match item.text_edit { + Some(CompletionItemTextEdit::TextEdit(edit)) => Some(edit.new_text), + _ => None, + }) + .collect(); + assert_eq!(new_texts, vec![r#"["say \"hi\""]"#, r#"["two words"]"#]); +} + +#[test] +fn typed_dict_attribute_completion_skips_non_typed_dict_bases() { + let code = r#" +from typing import TypedDict + +class A(TypedDict): + f1: int + +class C: + f1: int + +c = C() +c. +# ^ +A. +# ^ +"#; + let (handles, state) = mk_multi_file_state(&[("main", code)], Require::Exports, false); + for position in extract_cursors_for_test(code) { + let completions = state.transaction().completion( + &handles["main"], + position, + ImportFormat::Absolute, + true, + None, + ); + assert!(!completions.is_empty()); + assert!( + completions.iter().all(|item| item.text_edit.is_none()), + "{completions:?}" + ); + } +} + #[test] fn builtins_doesnt_autoimport() { let code = r#" From f23386e7c6ee124e437dca32ff068262625b8cd6 Mon Sep 17 00:00:00 2001 From: Parman Mohammadalizadeh Date: Sun, 1 Nov 2026 11:39:00 +0100 Subject: [PATCH 2/4] Offer TypedDict keys as completions after a dot TypedDict keys can only be read with a subscript, so after `a.` users saw the mapping methods but none of the keys they usually want (#5478). Attribute completion now also offers each key of a TypedDict base. The item carries a text edit that replaces the dot and any typed prefix with `["key"]`, so accepting `f1` after `a.` produces `a["f1"]`. The keys come from the same TypedDict field lookup that subscript key completion uses, so inherited, NotRequired and ReadOnly keys are all included, and the regular attribute completions are still added alongside them. The filter text starts with the text the edit replaces (the dot), since clients match it against what was typed from the start of the edit. Keys are quoted with JSON string escaping, which is also a valid Python string literal, so keys that are not identifiers stay correct. The items are skipped when the base and attribute are on different lines, because a completion text edit must stay on one line. --- pyrefly/lib/lsp/wasm/completion.rs | 9 ++++- pyrefly/lib/state/lsp/dict_completions.rs | 46 +++++++++++++++++++++++ 2 files changed, 54 insertions(+), 1 deletion(-) diff --git a/pyrefly/lib/lsp/wasm/completion.rs b/pyrefly/lib/lsp/wasm/completion.rs index 93980e4695..43c1e467eb 100644 --- a/pyrefly/lib/lsp/wasm/completion.rs +++ b/pyrefly/lib/lsp/wasm/completion.rs @@ -1190,7 +1190,7 @@ impl Transaction<'_> { } } Some(IdentifierWithContext { - identifier: _, + identifier, context: IdentifierContext::Attribute { base_range, .. }, }) => { let expected_type = covering_nodes @@ -1224,6 +1224,13 @@ impl Transaction<'_> { } } } + self.add_typed_dict_attribute_key_completions( + handle, + &base_type, + base_range, + &identifier, + &mut result, + ); self.add_attribute_completions_for_type( handle, base_type, diff --git a/pyrefly/lib/state/lsp/dict_completions.rs b/pyrefly/lib/state/lsp/dict_completions.rs index 11bd49f76b..5830ae0f29 100644 --- a/pyrefly/lib/state/lsp/dict_completions.rs +++ b/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; @@ -27,6 +29,7 @@ use ruff_python_ast::ModModule; 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; @@ -1090,6 +1093,49 @@ 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, + completions: &mut Vec, + ) { + if !Self::type_contains_typed_dict(base_type) { + return; + } + let Some(module_info) = self.get_module_info(handle) else { + return; + }; + let edit_range = + module_info.to_lsp_range(TextRange::new(base_range.end(), attribute.range().end())); + // LSP requires a completion's text edit to be on a single line. + if edit_range.start.line != edit_range.end.line { + 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(base_range.end(), attribute.start())); + for (key, ty) in keys { + let quoted_key = to_string(&key).expect("Serializing a string cannot fail"); + 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, From 4b272fc1de46c69af499322616a3d363212766e4 Mon Sep 17 00:00:00 2001 From: Parman Mohammadalizadeh Date: Mon, 2 Nov 2026 21:46:29 +0100 Subject: [PATCH 3/4] Add failing tests for review edge cases in TypedDict dot completion Review of #5480 found three places where accepting a TypedDict key after a dot produces invalid code. In `(a).fi` the edit swallows the closing parenthesis. Inside `f"{a.fi}"` the double-quoted key reuses the f-string's quote, which is a syntax error before Python 3.12. In `case a.fi:` the subscript is not a valid value pattern. These tests pin the expected behavior for each. --- pyrefly/lib/test/lsp/completion.rs | 86 ++++++++++++++++++++++++++++++ 1 file changed, 86 insertions(+) diff --git a/pyrefly/lib/test/lsp/completion.rs b/pyrefly/lib/test/lsp/completion.rs index 9b7862eda6..e6cb247d77 100644 --- a/pyrefly/lib/test/lsp/completion.rs +++ b/pyrefly/lib/test/lsp/completion.rs @@ -3087,6 +3087,92 @@ A. } } +#[test] +fn typed_dict_attribute_completion_keeps_parentheses_around_the_base() { + let code = r#" +from typing import TypedDict + +class A(TypedDict): + field: int + +a: A +(a).fi +# ^ +"#; + let completions = completions_at_cursor(code); + let field = completions + .iter() + .find(|item| item.label == "field") + .expect("missing completion for TypedDict key `field`"); + assert_eq!(field.filter_text.as_deref(), Some(".field")); + assert_eq!( + field.text_edit, + Some(CompletionItemTextEdit::TextEdit(TextEdit { + range: Range::new(Position::new(7, 3), Position::new(7, 6)), + new_text: r#"["field"]"#.to_owned(), + })) + ); +} + +#[test] +fn typed_dict_attribute_completion_avoids_enclosing_f_string_quotes() { + let code = r#" +from typing import TypedDict + +Keys = TypedDict("Keys", {"field": int, "fi'eld": int, 'fi"eld': int}) + +k: Keys +f"{k.fi}" +# ^ +f'{k.fi}' +# ^ +f"{f'{k.fi}'}" +# ^ +"#; + let (handles, state) = mk_multi_file_state(&[("main", code)], Require::Exports, false); + let new_texts: Vec> = extract_cursors_for_test(code) + .into_iter() + .map(|position| { + state + .transaction() + .completion( + &handles["main"], + position, + ImportFormat::Absolute, + true, + None, + ) + .into_iter() + .filter_map(|item| match item.text_edit { + Some(CompletionItemTextEdit::TextEdit(edit)) => Some(edit.new_text), + _ => None, + }) + .collect() + }) + .collect(); + assert_eq!( + new_texts, + vec![vec!["['field']"], vec![r#"["field"]"#], Vec::<&str>::new()] + ); +} + +#[test] +fn typed_dict_attribute_completion_skips_value_patterns() { + let code = r#" +from typing import TypedDict + +class A(TypedDict): + field: int + +a: A +match 0: + case a.fi: +# ^ + pass +"#; + assert!(subscript_rewrite_labels(code).is_empty()); +} + #[test] fn builtins_doesnt_autoimport() { let code = r#" From 7298ce8f8a4e1deeb4fb3cb2d430ea36cd830898 Mon Sep 17 00:00:00 2001 From: Parman Mohammadalizadeh Date: Mon, 2 Nov 2026 21:46:30 +0100 Subject: [PATCH 4/4] Keep TypedDict dot completion valid in parentheses, f-strings and patterns The subscript rewrite assumed the dot sits right after the base expression, that a double-quoted key is always a valid literal, and that a subscript is allowed wherever an attribute is. None of these hold everywhere. The base range excludes grouping parentheses, so the edit now starts at the dot itself, found in the source between the base and the attribute name. Inside an f-string the key is quoted with a quote character that no enclosing f-string uses, and keys that would need a quote or an escape are skipped, since before PEP 701 a replacement field cannot contain either. Value patterns only allow dotted names, so the items are not offered there at all. --- pyrefly/lib/lsp/wasm/completion.rs | 22 +++++++++---- pyrefly/lib/state/lsp/dict_completions.rs | 39 ++++++++++++++++++++--- 2 files changed, 50 insertions(+), 11 deletions(-) diff --git a/pyrefly/lib/lsp/wasm/completion.rs b/pyrefly/lib/lsp/wasm/completion.rs index 43c1e467eb..15bff086ce 100644 --- a/pyrefly/lib/lsp/wasm/completion.rs +++ b/pyrefly/lib/lsp/wasm/completion.rs @@ -1224,13 +1224,21 @@ impl Transaction<'_> { } } } - self.add_typed_dict_attribute_key_completions( - handle, - &base_type, - base_range, - &identifier, - &mut result, - ); + // 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, diff --git a/pyrefly/lib/state/lsp/dict_completions.rs b/pyrefly/lib/state/lsp/dict_completions.rs index 5830ae0f29..9eef71489f 100644 --- a/pyrefly/lib/state/lsp/dict_completions.rs +++ b/pyrefly/lib/state/lsp/dict_completions.rs @@ -26,6 +26,8 @@ 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; @@ -1101,6 +1103,7 @@ impl<'a> Transaction<'a> { base_type: &Type, base_range: TextRange, attribute: &Identifier, + covering_nodes: &[AnyNodeRef], completions: &mut Vec, ) { if !Self::type_contains_typed_dict(base_type) { @@ -1109,19 +1112,47 @@ impl<'a> Transaction<'a> { let Some(module_info) = self.get_module_info(handle) else { return; }; - let edit_range = - module_info.to_lsp_range(TextRange::new(base_range.end(), attribute.range().end())); + // 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 = 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(base_range.end(), attribute.start())); + let accessor = module_info.code_at(TextRange::new(dot, attribute.start())); for (key, ty) in keys { - let quoted_key = to_string(&key).expect("Serializing a string cannot fail"); + 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,