diff --git a/pyrefly/lib/lsp/wasm/completion.rs b/pyrefly/lib/lsp/wasm/completion.rs index 298ea4cb57..e3535397e3 100644 --- a/pyrefly/lib/lsp/wasm/completion.rs +++ b/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, diff --git a/pyrefly/lib/state/lsp/dict_completions.rs b/pyrefly/lib/state/lsp/dict_completions.rs index 2c1fd48350..9d4b72e8e4 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; @@ -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, + ) { + 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 = 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, diff --git a/pyrefly/lib/test/lsp/completion.rs b/pyrefly/lib/test/lsp/completion.rs index 2c98d1e562..dd732776d2 100644 --- a/pyrefly/lib/test/lsp/completion.rs +++ b/pyrefly/lib/test/lsp/completion.rs @@ -11,6 +11,9 @@ 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; @@ -3177,6 +3180,258 @@ 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 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#"