diff --git a/pyrefly/lib/lsp/wasm/completion.rs b/pyrefly/lib/lsp/wasm/completion.rs index e3535397e3..f536c47383 100644 --- a/pyrefly/lib/lsp/wasm/completion.rs +++ b/pyrefly/lib/lsp/wasm/completion.rs @@ -1283,12 +1283,63 @@ impl Transaction<'_> { if let Some(nodes) = covering_nodes.as_deref() { self.add_match_class_keyword_completions(handle, nodes, position, &mut result); } - // Because of parser error recovery, `from x impo...` looks like `from x import impo...` - // If the user might be typing the `import` keyword, add that as an autocomplete option. - match covering_nodes + let mut at_statement_keyword = false; + let identifier_with_context = covering_nodes .as_deref() .and_then(Self::identifier_from_covering_nodes) - { + .or_else(|| { + // Parser recovery can discard a keyword such as `as` at statement start. + // Recover its prefix so filtering and deferred auto-imports use the normal path. + let tokens = self.get_parsed_module(handle)?.tokens()?; + let keyword = tokens + .at_offset(position) + .find(|token| token.kind().is_keyword() && token.start() < position)?; + let previous = tokens.before(keyword.start()).iter().rev().find(|token| { + !token.kind().is_trivia() + && !matches!(token.kind(), TokenKind::Indent | TokenKind::Dedent) + }); + let at_statement_start = match previous { + None => true, + Some(token) if matches!(token.kind(), TokenKind::Newline | TokenKind::Semi) => { + true + } + Some(colon) if colon.kind() == TokenKind::Colon => { + // Only a compound-statement colon can introduce an inline suite. + // Exclude child nodes that end immediately before the colon. + matches!( + Ast::locate_node(ast.as_ref()?.as_ref(), colon.start()) + .into_iter() + .find(|node| node.range().contains(colon.start())), + Some( + AnyNodeRef::StmtFunctionDef(_) + | AnyNodeRef::StmtClassDef(_) + | AnyNodeRef::StmtIf(_) + | AnyNodeRef::StmtWhile(_) + | AnyNodeRef::StmtFor(_) + | AnyNodeRef::StmtWith(_) + | AnyNodeRef::StmtTry(_) + | AnyNodeRef::ElifElseClause(_) + | AnyNodeRef::ExceptHandlerExceptHandler(_) + | AnyNodeRef::MatchCase(_) + ) + ) + } + _ => false, + }; + if !at_statement_start { + return None; + } + let module_info = self.get_module_info(handle)?; + let range = TextRange::new(keyword.start(), position); + at_statement_keyword = true; + Some(IdentifierWithContext { + identifier: Identifier::new(Name::new(module_info.code_at(range)), range), + context: IdentifierContext::Expr(ExprContext::Load), + }) + }); + // Because of parser error recovery, `from x impo...` looks like `from x import impo...` + // If the user might be typing the `import` keyword, add that as an autocomplete option. + match identifier_with_context { Some(IdentifierWithContext { context: IdentifierContext::AliasDefinition, .. @@ -1477,10 +1528,11 @@ impl Transaction<'_> { .as_deref() .is_some_and(|nodes| Self::is_typing_keyword_argument_name(nodes, position)); if !skip_value_completions && !is_method_def { - let at_statement_start = matches!( - covering_nodes.as_deref().and_then(|nodes| nodes.get(1)), - Some(AnyNodeRef::StmtExpr(_)) - ); + let at_statement_start = at_statement_keyword + || matches!( + covering_nodes.as_deref().and_then(|nodes| nodes.get(1)), + Some(AnyNodeRef::StmtExpr(_)) + ); let expression_only = matches!(context, IdentifierContext::Expr(_)) && !at_statement_start; Self::add_keyword_completions(handle, expression_only, &mut result); diff --git a/pyrefly/lib/test/lsp/completion.rs b/pyrefly/lib/test/lsp/completion.rs index dd732776d2..080b150a55 100644 --- a/pyrefly/lib/test/lsp/completion.rs +++ b/pyrefly/lib/test/lsp/completion.rs @@ -19,6 +19,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; @@ -2715,6 +2716,157 @@ def f(): ); } +#[test] +fn completion_after_keyword_prefix_in_function_body() { + for (prefix, local, builtin) in [ + ("as", "assertion", "classmethod"), + ("in", "initial", "int"), + ("is", "island", "isinstance"), + ] { + let code = format!("\ndef f({local}: int, unrelated: int):\n {prefix}\n# ^\n"); + let items = completions_at_cursor(&code); + for keyword in ["assert", "async", "import", "is"] { + assert!( + items.iter().any(|item| { + item.kind == Some(CompletionItemKind::Keyword) && item.label == keyword + }), + "expected keyword `{keyword}` after `{prefix}`, got {items:?}" + ); + } + for name in [local, builtin] { + assert!( + items.iter().any(|item| item.label == name), + "expected `{name}` after `{prefix}`, got {items:?}" + ); + } + for name in ["unrelated", "len"] { + assert!( + !items.iter().any(|item| item.label == name), + "unexpected `{name}` after `{prefix}`, got {items:?}" + ); + } + } +} + +#[test] +fn completion_after_keyword_prefix_preserves_autoimport_behavior() { + for (prefix, local, auto_import, expected_incomplete, expected_autoimport) in [ + ("as", "", true, true, false), + ("ass", "", true, false, true), + ("and", "", true, false, true), + ("and", "andromeda: int", true, true, false), + ("and_", "andromeda: int", true, false, true), + ("as", "", false, false, false), + ("and", "andromeda: int", false, false, false), + ] { + let code = format!("def f({local}):\n {prefix}"); + let (handles, state) = mk_multi_file_state( + &[ + ("main", &code), + ("tools", "def and_tool(): pass\ndef assert_tool(): pass"), + ], + Require::Exports, + false, + ); + let (items, is_incomplete) = state.transaction().completion_with_incomplete( + &handles["main"], + TextSize::of(code.as_str()), + ImportFormat::Absolute, + CompletionOptions { + auto_import, + ..Default::default() + }, + None, + ); + assert!(!items.is_empty(), "expected completions in {code:?}"); + assert_eq!(is_incomplete, expected_incomplete, "code: {code:?}"); + assert_eq!( + items + .iter() + .any(|item| matches!(item.label.as_str(), "and_tool" | "assert_tool")), + expected_autoimport, + "code: {code:?}" + ); + } +} + +#[test] +fn completion_after_keyword_prefix_in_inline_suite() { + for header in [ + "def f():\n if True:", + "def f():\n while True:", + "def f():\n for item in ():", + "def f():\n with open('file'):", + "def f():\n try:", + "def f():\n try: pass\n except:", + "def f():\n try: pass\n finally:", + "def f():\n if False: pass\n elif True:", + "def f():\n if False: pass\n else:", + "class C:\n def f(self):", + "def f():\n class C:", + "def f():\n match 1:\n case 1:", + ] { + let code = format!("{header} as"); + let (handles, state) = mk_multi_file_state(&[("main", &code)], Require::Exports, false); + let items = state.transaction().completion( + &handles["main"], + TextSize::of(code.as_str()), + ImportFormat::Absolute, + true, + None, + ); + assert!( + items.iter().any(|item| { + item.kind == Some(CompletionItemKind::Keyword) && item.label == "assert" + }), + "expected `assert` in {code:?}, got {items:?}" + ); + } +} + +#[test] +fn completion_after_keyword_prefix_in_other_statement_positions() { + for code in [ + "as\n# ^\n", + "class C:\n as\n# ^\n", + "if True:\n as\n# ^\n", + "def f():\n if True:\n as\n# ^\n", + "def f():\n # comment\n\n as\n# ^\n", + "def f():\n pass; as\n# ^\n", + "def f():\r\n as\r\n# ^\r\n", + ] { + let items = completions_at_cursor(code); + assert!( + items.iter().any(|item| { + item.kind == Some(CompletionItemKind::Keyword) && item.label == "assert" + }), + "expected `assert` in {code:?}, got {items:?}" + ); + } +} + +#[test] +fn no_statement_completions_after_keyword_in_expression() { + for code in [ + "def f(item):\n value = item is\n# ^\n", + "def f():\n return as\n# ^\n", + "def f():\n print(\n as\n# ^\n )\n", + "def f():\n value = {1: as}\n# ^\n", + "def f(items):\n value = items[:as]\n# ^\n", + "def f():\n value: as\n# ^\n", + "def f():\n value = lambda: as\n# ^\n", + "def f():\n def g(value: as): pass\n# ^\n", + ] { + let items = completions_at_cursor(code); + for name in ["assert", "def", "return"] { + assert!( + !items.iter().any(|item| item.label == name), + "unexpected `{name}` in {code:?}, got {items:?}" + ); + } + } +} + #[test] fn kwargs_completion_with_existing_args() { let code = r#"