PY-20770 Fixed: Support Python 3.6 asynchronous generators and comprehensions

Update CompatibilityVisitor to highlight "async" inside comprehensions and generator expressions in Pythons < 3.6

python/src/com/jetbrains/python/validation/CompatibilityVisitor.java

@@ -25,6 +25,7 @@ import com.intellij.openapi.util.text.StringUtil;
 import com.intellij.psi.PsiElement;
 import com.intellij.psi.PsiWhiteSpace;
 import com.intellij.psi.tree.IElementType;
+import com.intellij.psi.tree.TokenSet;
 import com.intellij.psi.util.PsiTreeUtil;
 import com.intellij.psi.util.QualifiedName;
 import com.intellij.util.ArrayUtil;
@@ -717,4 +718,18 @@ public abstract class CompatibilityVisitor extends PyAnnotator {
       }
     }
   }
+
+  @Override
+  public void visitPyComprehensionElement(PyComprehensionElement node) {
+    super.visitPyComprehensionElement(node);
+
+    if (myVersionsToProcess.contains(LanguageLevel.PYTHON35)) {
+      Arrays
+        .stream(node.getNode().getChildren(TokenSet.create(PyTokenTypes.ASYNC_KEYWORD)))
+        .filter(Objects::nonNull)
+        .map(ASTNode::getPsi)
+        .forEach(element -> registerProblem(element,
+                                            "Python version 3.5 does not support 'async' inside comprehensions and generator expressions"));
+    }
+  }
 }

python/testData/highlighting/asyncComprehensionsPy35.py

@@ -0,0 +1,16 @@
+async def asyncgen():
+    <error descr="Python version 3.5 does not support 'yield' inside async functions">yield 10</error>
+async def run():
+    {i <error descr="Python version 3.5 does not support 'async' inside comprehensions and generator expressions">async</error> for i in asyncgen()}
+    [i <error descr="Python version 3.5 does not support 'async' inside comprehensions and generator expressions">async</error> for i in asyncgen()]
+    {i: i ** 2 <error descr="Python version 3.5 does not support 'async' inside comprehensions and generator expressions">async</error> for i in asyncgen()}
+    (i ** 2 <error descr="Python version 3.5 does not support 'async' inside comprehensions and generator expressions">async</error> for i in asyncgen())
+    list(i <error descr="Python version 3.5 does not support 'async' inside comprehensions and generator expressions">async</error> for i in asyncgen())
+
+    dataset = {data for line in gen()
+                    <error descr="Python version 3.5 does not support 'async' inside comprehensions and generator expressions">async</error> for data in line
+                    if check(data)}
+
+    dataset = {data <error descr="Python version 3.5 does not support 'async' inside comprehensions and generator expressions">async</error> for line in asyncgen()
+                    <error descr="Python version 3.5 does not support 'async' inside comprehensions and generator expressions">async</error> for data in line
+                    if check(data)}

python/testData/highlighting/asyncComprehensionsPy36.py

@@ -0,0 +1,16 @@
+async def asyncgen():
+    yield 10
+async def run():
+    {i async for i in asyncgen()}
+    [i async for i in asyncgen()]
+    {i: i ** 2 async for i in asyncgen()}
+    (i ** 2 async for i in asyncgen())
+    list(i async for i in asyncgen())
+
+    dataset = {data for line in gen()
+                    async for data in line
+                    if check(data)}
+
+    dataset = {data async for line in asyncgen()
+                    async for data in line
+                    if check(data)}

python/testData/inspections/PyCompatibilityInspection/asyncComprehensions.py

@@ -0,0 +1,16 @@
+<warning descr="Python versions < 3.5 do not support this syntax">async</warning> def asyncgen():
+    <warning descr="Python version 3.5 does not support 'yield' inside async functions">yield 10</warning>
+<warning descr="Python versions < 3.5 do not support this syntax">async</warning> def run():
+    <warning descr="Python version 2.4, 2.5, 2.6, 3.0 do not support set comprehensions">{i <warning descr="Python version 3.5 does not support 'async' inside comprehensions and generator expressions">async</warning> for i in asyncgen()}</warning>
+    [i <warning descr="Python version 3.5 does not support 'async' inside comprehensions and generator expressions">async</warning> for i in asyncgen()]
+    <warning descr="Python version 2.4, 2.5, 2.6, 3.0 do not support dictionary comprehensions">{i: i ** 2 <warning descr="Python version 3.5 does not support 'async' inside comprehensions and generator expressions">async</warning> for i in asyncgen()}</warning>
+    (i ** 2 <warning descr="Python version 3.5 does not support 'async' inside comprehensions and generator expressions">async</warning> for i in asyncgen())
+    list(i <warning descr="Python version 3.5 does not support 'async' inside comprehensions and generator expressions">async</warning> for i in asyncgen())
+
+    dataset = <warning descr="Python version 2.4, 2.5, 2.6, 3.0 do not support set comprehensions">{data for line in gen()
+                    <warning descr="Python version 3.5 does not support 'async' inside comprehensions and generator expressions">async</warning> for data in line
+                    if check(data)}</warning>
+
+    dataset = <warning descr="Python version 2.4, 2.5, 2.6, 3.0 do not support set comprehensions">{data <warning descr="Python version 3.5 does not support 'async' inside comprehensions and generator expressions">async</warning> for line in asyncgen()
+                    <warning descr="Python version 3.5 does not support 'async' inside comprehensions and generator expressions">async</warning> for data in line
+                    if check(data)}</warning>

python/testSrc/com/jetbrains/python/PythonHighlightingTest.java

@@ -310,6 +310,16 @@ public class PythonHighlightingTest extends PyTestCase {
     doTest();
   }
 
+  // PY-20770
+  public void testAsyncComprehensionsPy35() {
+    doTest(LanguageLevel.PYTHON35, true, false);
+  }
+
+  // PY-20770
+  public void testAsyncComprehensionsPy36() {
+    doTest(LanguageLevel.PYTHON36, true, false);
+  }
+
   // PY-20775
   public void testFStringMissingRightBrace() {
     runWithLanguageLevel(LanguageLevel.PYTHON36, () -> doTest(true, false));

python/testSrc/com/jetbrains/python/inspections/PyCompatibilityInspectionTest.java

@@ -199,6 +199,11 @@ public class PyCompatibilityInspectionTest extends PyTestCase {
     doTest(LanguageLevel.PYTHON36);
   }
 
+  // PY-20770
+  public void testAsyncComprehensions() {
+    doTest(LanguageLevel.PYTHON36);
+  }
+
   private void doTest(@NotNull LanguageLevel level) {
     runWithLanguageLevel(level, this::doTest);
   }