diff --git a/python/python-psi-impl/resources/messages/PyPsiBundle.properties b/python/python-psi-impl/resources/messages/PyPsiBundle.properties index 9e891092d84a..5615a961424f 100644 --- a/python/python-psi-impl/resources/messages/PyPsiBundle.properties +++ b/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 diff --git a/python/python-psi-impl/src/com/jetbrains/python/validation/PyAsyncAwaitAnnotator.java b/python/python-psi-impl/src/com/jetbrains/python/validation/PyAsyncAwaitAnnotator.java index d178fe1eea8e..b4e1dd7d446a 100644 --- a/python/python-psi-impl/src/com/jetbrains/python/validation/PyAsyncAwaitAnnotator.java +++ b/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 { protected ConvertIntoAsyncFunctionFix(@NotNull PyFunction element) { super(element); diff --git a/python/testData/highlighting/asyncDictComprehensionInNonAsyncFunction.after.py b/python/testData/highlighting/asyncDictComprehensionInNonAsyncFunction.after.py new file mode 100644 index 000000000000..c790ecdc134c --- /dev/null +++ b/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()} diff --git a/python/testData/highlighting/asyncDictComprehensionInNonAsyncFunction.py b/python/testData/highlighting/asyncDictComprehensionInNonAsyncFunction.py new file mode 100644 index 000000000000..dc141028d7b0 --- /dev/null +++ b/python/testData/highlighting/asyncDictComprehensionInNonAsyncFunction.py @@ -0,0 +1,10 @@ +import asyncio + +async def genfunc(): + yield 1 + +def example(): + {x: x async for x in genfunc()} + +async def example_correct(): + {x: x async for x in genfunc()} diff --git a/python/testData/highlighting/asyncForInNonAsyncFunction.after.py b/python/testData/highlighting/asyncForInNonAsyncFunction.after.py new file mode 100644 index 000000000000..c2f37df867ed --- /dev/null +++ b/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 diff --git a/python/testData/highlighting/asyncForInNonAsyncFunction.py b/python/testData/highlighting/asyncForInNonAsyncFunction.py new file mode 100644 index 000000000000..e09b8f9376d1 --- /dev/null +++ b/python/testData/highlighting/asyncForInNonAsyncFunction.py @@ -0,0 +1,12 @@ +import asyncio + +async def genfunc(): + yield 1 + +def example(): + async for x in genfunc(): + pass + +async def example_correct(): + async for x in genfunc(): + pass diff --git a/python/testData/highlighting/asyncListComprehensionInNonAsyncFunction.after.py b/python/testData/highlighting/asyncListComprehensionInNonAsyncFunction.after.py new file mode 100644 index 000000000000..5292d975777f --- /dev/null +++ b/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()] diff --git a/python/testData/highlighting/asyncListComprehensionInNonAsyncFunction.py b/python/testData/highlighting/asyncListComprehensionInNonAsyncFunction.py new file mode 100644 index 000000000000..cac023c24888 --- /dev/null +++ b/python/testData/highlighting/asyncListComprehensionInNonAsyncFunction.py @@ -0,0 +1,10 @@ +import asyncio + +async def genfunc(): + yield 1 + +def example(): + [x async for x in genfunc()] + +async def example_correct(): + [x async for x in genfunc()] diff --git a/python/testData/highlighting/asyncSetComprehensionInNonAsyncFunction.after.py b/python/testData/highlighting/asyncSetComprehensionInNonAsyncFunction.after.py new file mode 100644 index 000000000000..7d523ee9ad2d --- /dev/null +++ b/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()} diff --git a/python/testData/highlighting/asyncSetComprehensionInNonAsyncFunction.py b/python/testData/highlighting/asyncSetComprehensionInNonAsyncFunction.py new file mode 100644 index 000000000000..1ae03119e2a8 --- /dev/null +++ b/python/testData/highlighting/asyncSetComprehensionInNonAsyncFunction.py @@ -0,0 +1,10 @@ +import asyncio + +async def genfunc(): + yield 1 + +def example(): + {x async for x in genfunc()} + +async def example_correct(): + {x async for x in genfunc()} diff --git a/python/testData/highlighting/asyncWithInNonAsyncFunction.after.py b/python/testData/highlighting/asyncWithInNonAsyncFunction.after.py new file mode 100644 index 000000000000..6685840b78fb --- /dev/null +++ b/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 diff --git a/python/testData/highlighting/asyncWithInNonAsyncFunction.py b/python/testData/highlighting/asyncWithInNonAsyncFunction.py new file mode 100644 index 000000000000..552f2017e72e --- /dev/null +++ b/python/testData/highlighting/asyncWithInNonAsyncFunction.py @@ -0,0 +1,10 @@ +import asyncio +from contextlib import AsyncExitStack + +def example(): + async with AsyncExitStack(): + pass + +async def example_correct(): + async with AsyncExitStack(): + pass diff --git a/python/testData/highlighting/awaitInComprehensionInNonAsyncFunction.py b/python/testData/highlighting/awaitInComprehensionInNonAsyncFunction.py index b7c0d4d1b505..3918f216eff9 100644 --- a/python/testData/highlighting/awaitInComprehensionInNonAsyncFunction.py +++ b/python/testData/highlighting/awaitInComprehensionInNonAsyncFunction.py @@ -1,5 +1,5 @@ def example(y): - return [x async for x in await y] + return [x for x in await y] async def example_correct(y): - return [x async for x in await y] \ No newline at end of file + return [x for x in await y] \ No newline at end of file diff --git a/python/testData/highlighting/awaitInLoopInNonAsyncFunction.py b/python/testData/highlighting/awaitInLoopInNonAsyncFunction.py index 4919beab385a..4d57045c4198 100644 --- a/python/testData/highlighting/awaitInLoopInNonAsyncFunction.py +++ b/python/testData/highlighting/awaitInLoopInNonAsyncFunction.py @@ -1,7 +1,7 @@ def example(x): - async for i in await x: + for i in await x: yield i async def example_correct(x): - async for i in await x: + for i in await x: yield i diff --git a/python/testData/highlighting/awaitInNonAsyncFunction.after.py b/python/testData/highlighting/awaitInNonAsyncFunction.after.py index 8de3732de646..8bc1d27abbd3 100644 --- a/python/testData/highlighting/awaitInNonAsyncFunction.after.py +++ b/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) \ No newline at end of file + await asyncio.sleep(1) diff --git a/python/testData/highlighting/awaitInNonAsyncFunction.py b/python/testData/highlighting/awaitInNonAsyncFunction.py index 9c243b985cf9..ea20df356d87 100644 --- a/python/testData/highlighting/awaitInNonAsyncFunction.py +++ b/python/testData/highlighting/awaitInNonAsyncFunction.py @@ -1,6 +1,6 @@ import asyncio -def example(): +def await_example(): await asyncio.sleep(1) async def example_correct(): diff --git a/python/testSrc/com/jetbrains/python/Py3HighlightingTest.java b/python/testSrc/com/jetbrains/python/Py3HighlightingTest.java index 7d310f94c215..eb807ba3fd9b 100644 --- a/python/testSrc/com/jetbrains/python/Py3HighlightingTest.java +++ b/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 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); }