From d6a5bbd91c60ffb66327b378ca417824d757040b Mon Sep 17 00:00:00 2001
From: RasmusNygren <rasmus.nygren99@gmail.com>
Date: Mon, 15 Dec 2025 18:56:34 +0100
Subject: [PATCH] [ty] Remove invalid statement-keyword completions in
 for-statements (#21979)

In `for x in <CURSOR>` statements it's only valid to provide expressions
that eventually evaluate to an iterable. While it's extremely difficult
to know if something can evaulate to an iterable in a general case,
there are some suggestions we know can never lead to an iterable. Most
keywords are such and hence we remove them here.

## Summary
This suppresses statement-keywords from auto-complete suggestions in
`for x in <CURSOR>` statements where we know they can never be valid, as
whatever is typed has to (at some point) evaluate to an iterable.

It handles the core issue from
https://github.com/astral-sh/ty/issues/1774 but there's a lot of related
cases that probably has to be handled piece-wise.

## Test Plan
New tests and verifying in the playground.
---
 crates/ty_ide/src/completion.rs | 104 ++++++++++++++++++++++++++++++--
 1 file changed, 98 insertions(+), 6 deletions(-)

diff --git a/crates/ty_ide/src/completion.rs b/crates/ty_ide/src/completion.rs
index a33f7600eb..be1e0bfa4f 100644
--- a/crates/ty_ide/src/completion.rs
+++ b/crates/ty_ide/src/completion.rs
@@ -490,6 +490,17 @@ pub fn completion<'db>(
             !ty.is_notimplemented(db)
         });
     }
+    if is_specifying_for_statement_iterable(&parsed, offset, typed.as_deref()) {
+        // Remove all keywords that doesn't make sense given the context,
+        // even if they are syntatically valid, e.g. `None`.
+        completions.retain(|item| {
+            let Some(kind) = item.kind else { return true };
+            if kind != CompletionKind::Keyword {
+                return true;
+            }
+            matches!(item.name.as_str(), "await" | "lambda" | "yield")
+        });
+    }
     completions.into_completions()
 }
 
@@ -1607,12 +1618,7 @@ fn is_in_definition_place(
 /// Returns true when the cursor sits on a binding statement.
 /// E.g. naming a parameter, type parameter, or `for` <name>).
 fn is_in_variable_binding(parsed: &ParsedModuleRef, offset: TextSize, typed: Option<&str>) -> bool {
-    let range = if let Some(typed) = typed {
-        let start = offset.saturating_sub(typed.text_len());
-        TextRange::new(start, offset)
-    } else {
-        TextRange::empty(offset)
-    };
+    let range = typed_text_range(typed, offset);
 
     let covering = covering_node(parsed.syntax().into(), range);
     covering.ancestors().any(|node| match node {
@@ -1667,6 +1673,36 @@ fn is_raising_exception(tokens: &[Token]) -> bool {
     false
 }
 
+/// Returns true when the cursor is after the `in` keyword in a
+/// `for x in <CURSOR>` statement.
+fn is_specifying_for_statement_iterable(
+    parsed: &ParsedModuleRef,
+    offset: TextSize,
+    typed: Option<&str>,
+) -> bool {
+    let range = typed_text_range(typed, offset);
+
+    let covering = covering_node(parsed.syntax().into(), range);
+    covering.parent().is_some_and(|node| {
+        matches!(
+           node, ast::AnyNodeRef::StmtFor(stmt_for) if stmt_for.iter.range().contains_range(range)
+        )
+    })
+}
+
+/// Returns the `TextRange` of the `typed` text.
+///
+/// `typed` should be the text immediately before the
+/// provided cursor `offset`.
+fn typed_text_range(typed: Option<&str>, offset: TextSize) -> TextRange {
+    if let Some(typed) = typed {
+        let start = offset.saturating_sub(typed.text_len());
+        TextRange::new(start, offset)
+    } else {
+        TextRange::empty(offset)
+    }
+}
+
 /// Order completions according to the following rules:
 ///
 /// 1) Names with no underscore prefix
@@ -5865,6 +5901,62 @@ def foo(param: s<CURSOR>)
         .contains("str");
     }
 
+    #[test]
+    fn no_statement_keywords_in_for_statement_simple1() {
+        completion_test_builder(
+            "\
+for x in a<CURSOR>
+",
+        )
+        .build()
+        .contains("lambda")
+        .contains("await")
+        .not_contains("raise")
+        .not_contains("False");
+    }
+
+    #[test]
+    fn no_statement_keywords_in_for_statement_simple2() {
+        completion_test_builder(
+            "\
+for x, y, _ in a<CURSOR>
+",
+        )
+        .build()
+        .contains("lambda")
+        .contains("await")
+        .not_contains("raise")
+        .not_contains("False");
+    }
+
+    #[test]
+    fn no_statement_keywords_in_for_statement_simple3() {
+        completion_test_builder(
+            "\
+for i, (x, y, z) in a<CURSOR>
+",
+        )
+        .build()
+        .contains("lambda")
+        .contains("await")
+        .not_contains("raise")
+        .not_contains("False");
+    }
+
+    #[test]
+    fn no_statement_keywords_in_for_statement_complex() {
+        completion_test_builder(
+            "\
+for i, (obj.x, (a[0], b['k']), _), *rest in a<CURSOR>
+",
+        )
+        .build()
+        .contains("lambda")
+        .contains("await")
+        .not_contains("raise")
+        .not_contains("False");
+    }
+
     #[test]
     fn favour_symbols_currently_imported() {
         let snapshot = CursorTest::builder()