PY-79522: Highlight syntax errors for misplaced "async with", "async for" and async list/set/dict comprehensions

# Conflicts:
#	community/python/python-psi-impl/src/com/jetbrains/python/validation/PyAsyncAwaitAnnotator.java

GitOrigin-RevId: 125a714f769ac93c8516250ae5b92afff14518f0

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

@@ -162,6 +162,8 @@ ANN.continue.break.or.return.in.star.except='break', 'continue' and 'return' can
 
 # PyAsyncAwaitAnnotator
 ANN.await.outside.async.function='await' outside async function
+ANN.async.with.outside.function='async with' outside async function
+ANN.async.for.outside.function='async for' outside async function
 QFIX.convert.into.async.function=Convert to async function
 
 ### quick doc generator

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

@@ -1,44 +1,104 @@
 package com.jetbrains.python.validation;
 
+import com.intellij.codeInspection.util.InspectionMessage;
 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.intellij.psi.PsiElement;
 import com.jetbrains.python.PyPsiBundle;
 import com.jetbrains.python.PyTokenTypes;
 import com.jetbrains.python.PythonRuntimeService;
+import com.jetbrains.python.codeInsight.controlflow.ScopeOwner;
 import com.jetbrains.python.codeInsight.dataflow.scope.ScopeUtil;
-import com.jetbrains.python.psi.PyExpressionCodeFragment;
-import com.jetbrains.python.psi.PyFunction;
-import com.jetbrains.python.psi.PyPrefixExpression;
+import com.jetbrains.python.psi.*;
 import org.jetbrains.annotations.NotNull;
 
 public final class PyAsyncAwaitAnnotator extends PyAnnotator {
+  private static boolean isAsyncAllowed(ScopeOwner scopeOwner) {
+    // Async functions are allowed to contain "await", "async with" and "async for"
+    if (scopeOwner instanceof PyFunction pyFunction && pyFunction.isAsync()) return true;
+
+    // Top-level expressions in the Python console are allowed to contain "await", "async with" and "async for"
+    if (scopeOwner instanceof PyExpressionCodeFragment && PythonRuntimeService.getInstance().isInPydevConsole(scopeOwner)) return true;
+
+    return false;
+  }
+
+
+  private void createError(@NotNull PsiElement node, ScopeOwner scopeOwner, @InspectionMessage @NotNull String message) {
+    var annotation = getHolder()
+      .newAnnotation(HighlightSeverity.ERROR, message)
+      .range(node);
+    if (scopeOwner instanceof PyFunction pyFunction) {
+      annotation = annotation.newFix(new ConvertIntoAsyncFunctionFix(pyFunction)).registerFix();
+    }
+    annotation.create();
+  }
+
+  private void checkComprehension(@NotNull PyComprehensionElement node) {
+    var asyncNode = node.getNode().findChildByType(PyTokenTypes.ASYNC_KEYWORD);
+    if (asyncNode == null) return;
+
+    var scopeOwner = ScopeUtil.getScopeOwner(node);
+    if (isAsyncAllowed(scopeOwner)) return;
+
+    createError((PsiElement)asyncNode, scopeOwner, PyPsiBundle.message("ANN.async.for.outside.function"));
+  }
+
   @Override
   public void visitPyPrefixExpression(@NotNull PyPrefixExpression node) {
     super.visitPyPrefixExpression(node);
+
     if (node.getOperator() == PyTokenTypes.AWAIT_KEYWORD) {
       var scopeOwner = ScopeUtil.getScopeOwner(node);
+      if (isAsyncAllowed(scopeOwner)) return;
 
-      // Async functions are allowed to contain "await"
-      if (scopeOwner instanceof PyFunction pyFunction && pyFunction.isAsync())
-        return;
-
-      // Top-level expressions in the Python console are allowed to contain "await"
-      if (scopeOwner instanceof PyExpressionCodeFragment && PythonRuntimeService.getInstance().isInPydevConsole(node))
-          return;
-
-      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();
+      createError(node.getFirstChild(), scopeOwner, PyPsiBundle.message("ANN.await.outside.async.function"));
     }
   }
 
+  @Override
+  public void visitPyForStatement(@NotNull PyForStatement node) {
+    super.visitPyForStatement(node);
+    if (!node.isAsync()) return;
+
+    var scopeOwner = ScopeUtil.getScopeOwner(node);
+    if (isAsyncAllowed(scopeOwner)) return;
+
+    createError(node.getFirstChild(), scopeOwner, PyPsiBundle.message("ANN.async.for.outside.function"));
+  }
+
+  @Override
+  public void visitPyWithStatement(@NotNull PyWithStatement node) {
+    super.visitPyWithStatement(node);
+    if (!node.isAsync()) return;
+
+    var scopeOwner = ScopeUtil.getScopeOwner(node);
+    if (isAsyncAllowed(scopeOwner)) return;
+
+    createError(node.getFirstChild(), scopeOwner, PyPsiBundle.message("ANN.async.with.outside.function"));
+  }
+
+  @Override
+  public void visitPyListCompExpression(@NotNull PyListCompExpression node) {
+    super.visitPyListCompExpression(node);
+    checkComprehension(node);
+  }
+
+  @Override
+  public void visitPyDictCompExpression(@NotNull PyDictCompExpression node) {
+    super.visitPyDictCompExpression(node);
+    checkComprehension(node);
+  }
+
+  @Override
+  public void visitPySetCompExpression(@NotNull PySetCompExpression node) {
+    super.visitPySetCompExpression(node);
+    checkComprehension(node);
+  }
+
   private static class ConvertIntoAsyncFunctionFix extends PsiUpdateModCommandAction<PyFunction> {
     protected ConvertIntoAsyncFunctionFix(@NotNull PyFunction element) {
       super(element);

python/testData/highlighting/asyncDictComprehensionInNonAsyncFunction.after.py

@@ -0,0 +1,10 @@
+import asyncio
+
+async def genfunc():
+    yield 1
+
+async def example():
+    {x: x async for x in genfunc()}
+
+async def example_correct():
+    {x: x async for x in genfunc()}

python/testData/highlighting/asyncDictComprehensionInNonAsyncFunction.py

@@ -0,0 +1,10 @@
+import asyncio
+
+async def genfunc():
+    yield 1
+
+def example():
+    {x: x <error descr="'async for' outside async function">async<caret></error> for x in genfunc()}
+
+async def example_correct():
+    {x: x async for x in genfunc()}

python/testData/highlighting/asyncForInNonAsyncFunction.after.py

@@ -0,0 +1,12 @@
+import asyncio
+
+async def genfunc():
+    yield 1
+
+async def example():
+    async for x in genfunc():
+        pass
+
+async def example_correct():
+    async for x in genfunc():
+        pass

python/testData/highlighting/asyncForInNonAsyncFunction.py

@@ -0,0 +1,12 @@
+import asyncio
+
+async def genfunc():
+    yield 1
+
+def example():
+    <error descr="'async for' outside async function">async<caret></error> for x in genfunc():
+        pass
+
+async def example_correct():
+    async for x in genfunc():
+        pass

python/testData/highlighting/asyncListComprehensionInNonAsyncFunction.after.py

@@ -0,0 +1,10 @@
+import asyncio
+
+async def genfunc():
+    yield 1
+
+async def example():
+    [x async for x in genfunc()]
+
+async def example_correct():
+    [x async for x in genfunc()]

python/testData/highlighting/asyncListComprehensionInNonAsyncFunction.py

@@ -0,0 +1,10 @@
+import asyncio
+
+async def genfunc():
+    yield 1
+
+def example():
+    [x <error descr="'async for' outside async function">async<caret></error> for x in genfunc()]
+
+async def example_correct():
+    [x async for x in genfunc()]

python/testData/highlighting/asyncSetComprehensionInNonAsyncFunction.after.py

@@ -0,0 +1,10 @@
+import asyncio
+
+async def genfunc():
+    yield 1
+
+async def example():
+    {x async for x in genfunc()}
+
+async def example_correct():
+    {x async for x in genfunc()}

python/testData/highlighting/asyncSetComprehensionInNonAsyncFunction.py

@@ -0,0 +1,10 @@
+import asyncio
+
+async def genfunc():
+    yield 1
+
+def example():
+    {x <error descr="'async for' outside async function">async<caret></error> for x in genfunc()}
+
+async def example_correct():
+    {x async for x in genfunc()}

python/testData/highlighting/asyncWithInNonAsyncFunction.after.py

@@ -0,0 +1,10 @@
+import asyncio
+from contextlib import AsyncExitStack
+
+async def example():
+    async with AsyncExitStack():
+        pass
+
+async def example_correct():
+    async with AsyncExitStack():
+        pass

python/testData/highlighting/asyncWithInNonAsyncFunction.py

@@ -0,0 +1,10 @@
+import asyncio
+from contextlib import AsyncExitStack
+
+def example():
+    <error descr="'async with' outside async function">async<caret></error> with AsyncExitStack():
+        pass
+
+async def example_correct():
+    async with AsyncExitStack():
+        pass

python/testData/highlighting/awaitInComprehensionInNonAsyncFunction.py

@@ -1,5 +1,5 @@
 def example(y):
-    return [x async for x in <error descr="'await' outside async function">await</error> y]
+    return [x for x in <error descr="'await' outside async function">await</error> y]
 
 async def example_correct(y):
-    return [x async for x in await y]
+    return [x for x in await y]

python/testData/highlighting/awaitInLoopInNonAsyncFunction.py

@@ -1,7 +1,7 @@
 def example(x):
-    async for i in <error descr="'await' outside async function">await</error> x:
+    for i in <error descr="'await' outside async function">await</error> x:
         yield i
 
 async def example_correct(x):
-    async for i in await x:
+    for i in await x:
         yield i

python/testData/highlighting/awaitInNonAsyncFunction.after.py

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

python/testData/highlighting/awaitInNonAsyncFunction.py

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

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

@@ -15,10 +15,8 @@
  */
 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 {
@@ -123,12 +121,32 @@ public class Py3HighlightingTest extends PyTestCase {
 
   // PY-32067
   public void testAwaitInNonAsyncFunction() {
-    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");
+    doHighlightingQuickfixTest("Convert to async function");
+  }
+
+  // PY-79522
+  public void testAsyncWithInNonAsyncFunction() {
+    doHighlightingQuickfixTest("Convert to async function");
+  }
+
+  // PY-79522
+  public void testAsyncForInNonAsyncFunction() {
+    doHighlightingQuickfixTest("Convert to async function");
+  }
+
+  // PY-79522
+  public void testAsyncListComprehensionInNonAsyncFunction() {
+    doHighlightingQuickfixTest("Convert to async function");
+  }
+
+  // PY-79522
+  public void testAsyncDictComprehensionInNonAsyncFunction() {
+    doHighlightingQuickfixTest("Convert to async function");
+  }
+
+  // PY-79522
+  public void testAsyncSetComprehensionInNonAsyncFunction() {
+    doHighlightingQuickfixTest("Convert to async function");
   }
 
   // PY-32067
@@ -186,6 +204,15 @@ public class Py3HighlightingTest extends PyTestCase {
     runWithLanguageLevel(languageLevel, () -> doTest(checkWarnings, checkInfos));
   }
 
+  private void doHighlightingQuickfixTest(String hint) {
+    var testPath = TEST_PATH + getTestName(true) + PyNames.DOT_PY;
+    var testPathAfter = TEST_PATH + getTestName(true) + ".after.py";
+    myFixture.testHighlighting(true, false, false, testPath);
+    var quickFix = myFixture.findSingleIntention(hint);
+    myFixture.launchAction(quickFix);
+    myFixture.testHighlighting(true, false, false, testPathAfter);
+  }
+
   private void doTest(boolean checkWarnings, boolean checkInfos) {
     myFixture.testHighlighting(checkWarnings, checkInfos, false, TEST_PATH + getTestName(true) + PyNames.DOT_PY);
   }