PY-32478: Suggest adding "async" to a function with the error "'await' outside async function"

Merge-request: IJ-MR-156125
Merged-by: Andrey Vokin <andrey.vokin@jetbrains.com>

(cherry picked from commit f3a2bc86331013879716b6991da26307c63b10fc)

IJ-MR-156125

GitOrigin-RevId: 10f0b3a8f60469af64bfd97613f258a282c9d0e2

python/python-psi-impl/resources/messages/PyPsiBundle.properties

@@ -160,7 +160,9 @@ ANN.try.except.can.not.have.except.and.star.except=Try statement cannot contain
 ANN.exception.group.in.star.except=ExceptionGroup cannot be used in except*
 ANN.continue.break.or.return.in.star.except='break', 'continue' and 'return' cannot appear in an except* block
 
+# PyAsyncAwaitAnnotator
 ANN.await.outside.async.function='await' outside async function
+QFIX.convert.into.async.function=Convert to async function
 
 ### quick doc generator
 QDOC.module.path.unknown=(Module path is unknown)

python/python-psi-impl/src/com/jetbrains/python/validation/PyAsyncAwaitAnnotator.java

@@ -1,5 +1,10 @@
 package com.jetbrains.python.validation;
 
+import com.intellij.lang.ASTNode;
+import com.intellij.lang.annotation.HighlightSeverity;
+import com.intellij.modcommand.ActionContext;
+import com.intellij.modcommand.ModPsiUpdater;
+import com.intellij.modcommand.PsiUpdateModCommandAction;
 import com.jetbrains.python.PyPsiBundle;
 import com.jetbrains.python.PyTokenTypes;
 import com.jetbrains.python.codeInsight.dataflow.scope.ScopeUtil;
@@ -14,8 +19,31 @@ public final class PyAsyncAwaitAnnotator extends PyAnnotator {
     if (node.getOperator() == PyTokenTypes.AWAIT_KEYWORD) {
       var scopeOwner = ScopeUtil.getScopeOwner(node);
       if (!(scopeOwner instanceof PyFunction pyFunction && pyFunction.isAsync())) {
-        markError(node.getFirstChild(), PyPsiBundle.message("ANN.await.outside.async.function"));
+        var annotation = getHolder()
+          .newAnnotation(HighlightSeverity.ERROR, PyPsiBundle.message("ANN.await.outside.async.function"))
+          .range(node.getFirstChild());
+        if (scopeOwner instanceof PyFunction pyFunction) {
+            annotation = annotation.newFix(new ConvertIntoAsyncFunctionFix(pyFunction)).registerFix();
+        }
+        annotation.create();
       }
     }
   }
+
+  private static class ConvertIntoAsyncFunctionFix extends PsiUpdateModCommandAction<PyFunction> {
+    protected ConvertIntoAsyncFunctionFix(@NotNull PyFunction element) {
+      super(element);
+    }
+
+    @Override
+    protected void invoke(@NotNull ActionContext context, @NotNull PyFunction element, @NotNull ModPsiUpdater updater) {
+      ASTNode defKeyword = element.getNode().findChildByType(PyTokenTypes.DEF_KEYWORD);
+      element.getNode().addLeaf(PyTokenTypes.ASYNC_KEYWORD, "async ", defKeyword);
+    }
+
+    @Override
+    public @NotNull String getFamilyName() {
+      return PyPsiBundle.message("QFIX.convert.into.async.function");
+    }
+  }
 }

python/testData/highlighting/awaitInNonAsyncFunction.after.py

@@ -0,0 +1,7 @@
+import asyncio
+
+async def example():
+    await asyncio.sleep(1)
+
+async def example_correct():
+    await asyncio.sleep(1)

python/testData/highlighting/awaitInNonAsyncFunction.py

@@ -1,7 +1,7 @@
 import asyncio
 
 def example():
-    <error descr="'await' outside async function">await</error> asyncio.sleep(1)
+    <error descr="'await' outside async function">await<caret></error> asyncio.sleep(1)
 
 async def example_correct():
     await asyncio.sleep(1)

python/testSrc/com/jetbrains/python/Py3HighlightingTest.java

@@ -15,8 +15,10 @@
  */
 package com.jetbrains.python;
 
+import com.intellij.codeInsight.intention.IntentionAction;
 import com.jetbrains.python.fixtures.PyTestCase;
 import com.jetbrains.python.psi.LanguageLevel;
+import java.util.List;
 import org.jetbrains.annotations.NotNull;
 
 public class Py3HighlightingTest extends PyTestCase {
@@ -121,7 +123,12 @@ public class Py3HighlightingTest extends PyTestCase {
 
   // PY-32067
   public void testAwaitInNonAsyncFunction() {
-    doTest(false, false);
+    var testPath = TEST_PATH + getTestName(true) + PyNames.DOT_PY;
+    myFixture.testHighlighting(true, false, false, testPath);
+    final List<IntentionAction> quickFixes = myFixture.filterAvailableIntentions("Convert");
+    assertOneElement(quickFixes);
+    myFixture.launchAction(quickFixes.get(0));
+    myFixture.checkResultByFile(TEST_PATH + getTestName(true) + ".after.py");
   }
 
   // PY-32067