[ty] Don't provide completions when in class or function definition (#21146)

crates/ty_ide/src/completion.rs

@@ -212,7 +212,10 @@ pub fn completion<'db>(
     offset: TextSize,
 ) -> Vec<Completion<'db>> {
     let parsed = parsed_module(db, file).load(db);
-    if is_in_comment(&parsed, offset) || is_in_string(&parsed, offset) {
+
+    let tokens = tokens_start_before(parsed.tokens(), offset);
+
+    if is_in_comment(tokens) || is_in_string(tokens) || is_in_definition_place(db, tokens, file) {
         return vec![];
     }
 
@@ -829,8 +832,7 @@ fn find_typed_text(
 
 /// Whether the given offset within the parsed module is within
 /// a comment or not.
-fn is_in_comment(parsed: &ParsedModuleRef, offset: TextSize) -> bool {
+fn is_in_comment(tokens: &[Token]) -> bool {
-    let tokens = tokens_start_before(parsed.tokens(), offset);
     tokens.last().is_some_and(|t| t.kind().is_comment())
 }
 
@@ -839,8 +841,7 @@ fn is_in_comment(parsed: &ParsedModuleRef, offset: TextSize) -> bool {
 ///
 /// Note that this will return `false` when positioned within an
 /// interpolation block in an f-string or a t-string.
-fn is_in_string(parsed: &ParsedModuleRef, offset: TextSize) -> bool {
+fn is_in_string(tokens: &[Token]) -> bool {
-    let tokens = tokens_start_before(parsed.tokens(), offset);
     tokens.last().is_some_and(|t| {
         matches!(
             t.kind(),
@@ -849,6 +850,29 @@ fn is_in_string(parsed: &ParsedModuleRef, offset: TextSize) -> bool {
     })
 }
 
+/// If the tokens end with `class f` or `def f` we return true.
+/// If the tokens end with `class` or `def`, we return false.
+/// This is fine because we don't provide completions anyway.
+fn is_in_definition_place(db: &dyn Db, tokens: &[Token], file: File) -> bool {
+    tokens
+        .len()
+        .checked_sub(2)
+        .and_then(|i| tokens.get(i))
+        .is_some_and(|t| {
+            if matches!(
+                t.kind(),
+                TokenKind::Def | TokenKind::Class | TokenKind::Type
+            ) {
+                true
+            } else if t.kind() == TokenKind::Name {
+                let source = source_text(db, file);
+                &source[t.range()] == "type"
+            } else {
+                false
+            }
+        })
+}
+
 /// Order completions according to the following rules:
 ///
 /// 1) Names with no underscore prefix
@@ -4058,6 +4082,83 @@ def f[T](x: T):
         test.build().contains("__repr__");
     }
 
+    #[test]
+    fn no_completions_in_function_def_name() {
+        let builder = completion_test_builder(
+            "\
+def f<CURSOR>
+    ",
+        );
+
+        builder.auto_import().build().not_contains("fabs");
+    }
+
+    #[test]
+    fn no_completions_in_function_def_empty_name() {
+        let builder = completion_test_builder(
+            "\
+def <CURSOR>
+        ",
+        );
+
+        builder.auto_import().build().not_contains("fabs");
+    }
+
+    #[test]
+    fn no_completions_in_class_def_name() {
+        let builder = completion_test_builder(
+            "\
+class f<CURSOR>
+    ",
+        );
+
+        builder.auto_import().build().not_contains("fabs");
+    }
+
+    #[test]
+    fn no_completions_in_class_def_empty_name() {
+        let builder = completion_test_builder(
+            "\
+class <CURSOR>
+        ",
+        );
+
+        builder.auto_import().build().not_contains("fabs");
+    }
+
+    #[test]
+    fn no_completions_in_type_def_name() {
+        let builder = completion_test_builder(
+            "\
+type f<CURSOR> = int
+    ",
+        );
+
+        builder.auto_import().build().not_contains("fabs");
+    }
+
+    #[test]
+    fn no_completions_in_maybe_type_def_name() {
+        let builder = completion_test_builder(
+            "\
+type f<CURSOR>
+       ",
+        );
+
+        builder.auto_import().build().not_contains("fabs");
+    }
+
+    #[test]
+    fn no_completions_in_type_def_empty_name() {
+        let builder = completion_test_builder(
+            "\
+type <CURSOR>
+        ",
+        );
+
+        builder.auto_import().build().not_contains("fabs");
+    }
+
     /// A way to create a simple single-file (named `main.py`) completion test
     /// builder.
     ///