fix Completion returns no items inside a function body when the typed text is a keyword (as, in, is, ...) #5529

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);

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#"