From 2e3d1fb6bc8f438cf94a4294b739f0b0a0a78d4e Mon Sep 17 00:00:00 2001
From: Azim Akhmadjonov <azimjon.akhmadjonov@jetbrains.com>
Date: Mon, 8 Dec 2025 10:13:52 +0000
Subject: [PATCH] PY-46932 add completion for "async def/for/with" constructs

Merge-request: IJ-MR-183838
Merged-by: Azimjon Akhmadjonov <azimjon.akhmadjonov@jetbrains.com>

GitOrigin-RevId: 49e89dd8faede246674cef69545c4c12f7128c3c
---
 .../PyKeywordCompletionContributor.java       | 33 +++++++++++++++++++
 python/testData/completion/async.after.py     |  1 -
 python/testData/completion/async.py           |  1 -
 .../jetbrains/python/Py3CompletionTest.java   |  4 ++-
 .../python/PythonKeywordCompletionTest.java   | 19 +++++++++++
 5 files changed, 55 insertions(+), 3 deletions(-)
 delete mode 100644 python/testData/completion/async.after.py
 delete mode 100644 python/testData/completion/async.py

diff --git a/python/python-psi-impl/src/com/jetbrains/python/codeInsight/completion/PyKeywordCompletionContributor.java b/python/python-psi-impl/src/com/jetbrains/python/codeInsight/completion/PyKeywordCompletionContributor.java
index 42f146819031..3ed838f70194 100644
--- a/python/python-psi-impl/src/com/jetbrains/python/codeInsight/completion/PyKeywordCompletionContributor.java
+++ b/python/python-psi-impl/src/com/jetbrains/python/codeInsight/completion/PyKeywordCompletionContributor.java
@@ -399,6 +399,14 @@ public final class PyKeywordCompletionContributor extends CompletionContributor
     result.addElement(TailTypeDecorator.withTail(lookupElement, tail));
   }
 
+  private static final InsertHandler<PythonLookupElement> TRAILING_SPACE_INSERT_HANDLER =
+    new InsertHandler<>() {
+      @Override
+      public void handleInsert(@NotNull InsertionContext context, @NotNull PythonLookupElement item) {
+        TailTypes.spaceType().processTail(context.getEditor(), context.getTailOffset());
+      }
+    };
+
   private void addPreColonStatements() {
     PsiElementPattern.Capture<PsiElement> startOfLine = psiElement()
       .withLanguage(PythonLanguage.getInstance())
@@ -685,6 +693,31 @@ public final class PyKeywordCompletionContributor extends CompletionContributor
   }
 
   private void addAsync() {
+    extend(CompletionType.BASIC,
+           psiElement()
+             .withLanguage(PythonLanguage.getInstance())
+             .and(PY35)
+             .and(IN_BEGIN_STMT)
+             .andNot(IN_COMMENT)
+             .andNot(IN_IMPORT_STMT)
+             .andNot(IN_PARAM_LIST)
+             .andNot(AFTER_QUALIFIER)
+             .andNot(IN_STRING_LITERAL)
+             .andNot(TARGET_AFTER_QUALIFIER)
+             .andNot(IN_PATTERN),
+           new CompletionProvider<>() {
+             @Override
+             protected void addCompletions(
+               @NotNull CompletionParameters parameters,
+               @NotNull ProcessingContext context,
+               @NotNull CompletionResultSet result
+             ) {
+               putKeyword(PyNames.ASYNC + " " + PyNames.DEF, TRAILING_SPACE_INSERT_HANDLER, TailTypes.noneType(), result);
+               putKeyword(PyNames.ASYNC + " " + PyNames.FOR, TRAILING_SPACE_INSERT_HANDLER, TailTypes.noneType(), result);
+               putKeyword(PyNames.ASYNC + " " + PyNames.WITH, TRAILING_SPACE_INSERT_HANDLER, TailTypes.noneType(), result);
+             }
+           });
+
     extend(CompletionType.BASIC,
            psiElement()
              .withLanguage(PythonLanguage.getInstance())
diff --git a/python/testData/completion/async.after.py b/python/testData/completion/async.after.py
deleted file mode 100644
index 500848fbe1b1..000000000000
--- a/python/testData/completion/async.after.py
+++ /dev/null
@@ -1 +0,0 @@
-async # Comment
diff --git a/python/testData/completion/async.py b/python/testData/completion/async.py
deleted file mode 100644
index ff8fb3181275..000000000000
--- a/python/testData/completion/async.py
+++ /dev/null
@@ -1 +0,0 @@
-asy<caret> # Comment
diff --git a/python/testSrc/com/jetbrains/python/Py3CompletionTest.java b/python/testSrc/com/jetbrains/python/Py3CompletionTest.java
index f6adb4dd4292..39ff03d11d5e 100644
--- a/python/testSrc/com/jetbrains/python/Py3CompletionTest.java
+++ b/python/testSrc/com/jetbrains/python/Py3CompletionTest.java
@@ -183,7 +183,9 @@ public class Py3CompletionTest extends PyTestCase {
   }
 
   public void testAsync() {
-    doTest();
+    List<String> suggested = doTestByText("asy<caret>");
+    assertNotNull(suggested);
+    assertContainsElements(suggested, "async", "async def", "async with", "async for");
   }
 
   public void testAwait() {
diff --git a/python/testSrc/com/jetbrains/python/PythonKeywordCompletionTest.java b/python/testSrc/com/jetbrains/python/PythonKeywordCompletionTest.java
index 0c23a7a71d3f..2da4bd7b7c00 100644
--- a/python/testSrc/com/jetbrains/python/PythonKeywordCompletionTest.java
+++ b/python/testSrc/com/jetbrains/python/PythonKeywordCompletionTest.java
@@ -360,4 +360,23 @@ public class PythonKeywordCompletionTest extends PyTestCase {
       assertContainsElements(variants, PyNames.DEF, PyNames.WITH, PyNames.FOR);
       assertDoesntContain(variants, "some_variable", "some_function", "SomeClass");
   }
+
+  // PY-46932
+  public void testAsyncKeywordConstructCompletion() {
+    List<String> variants = doTestByText("""
+                   some_variable = 42
+                   def some_function():
+                       pass
+                   class SomeClass:
+                       pass
+                   as<caret>
+                   """);
+
+    assertContainsElements(
+      variants,
+      PyNames.ASYNC + " " + PyNames.DEF,
+      PyNames.ASYNC + " " + PyNames.FOR,
+      PyNames.ASYNC + " " + PyNames.WITH
+    );
+  }
 }