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

Update highlighting and compatibility visitor to correctly process async generators described in PEP 525

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

@@ -31,6 +31,7 @@ import com.intellij.util.ArrayUtil;
 import com.intellij.util.containers.ContainerUtil;
 import com.jetbrains.python.PyNames;
 import com.jetbrains.python.PyTokenTypes;
+import com.jetbrains.python.codeInsight.dataflow.scope.ScopeUtil;
 import com.jetbrains.python.inspections.quickfix.*;
 import com.jetbrains.python.psi.*;
 import com.jetbrains.python.psi.impl.PyStringLiteralExpressionImpl;
@@ -449,6 +450,18 @@ public abstract class CompatibilityVisitor extends PyAnnotator {
   public void visitPyYieldExpression(PyYieldExpression node) {
     super.visitPyYieldExpression(node);
 
+    Optional
+      .ofNullable(ScopeUtil.getScopeOwner(node))
+      .map(owner -> PyUtil.as(owner, PyFunction.class))
+      .filter(function -> function.isAsync() && function.isAsyncAllowed())
+      .ifPresent(
+        function -> {
+          if (!node.isDelegating() && myVersionsToProcess.contains(LanguageLevel.PYTHON35)) {
+            registerProblem(node, "Python version 3.5 does not support 'yield' inside async functions");
+          }
+        }
+      );
+
     if (!node.isDelegating()) {
       return;
     }

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

@@ -1,5 +1,5 @@
 /*
- * Copyright 2000-2014 JetBrains s.r.o.
+ * Copyright 2000-2016 JetBrains s.r.o.
  *
  * Licensed under the Apache License, Version 2.0 (the "License");
  * you may not use this file except in compliance with the License.
@@ -21,7 +21,7 @@ import com.jetbrains.python.codeInsight.dataflow.scope.ScopeUtil;
 import com.jetbrains.python.psi.*;
 
 /**
- * Highlights incorrect return statements: 'return' and 'yield' outside functions, 'yield' inside async functions.
+ * Highlights incorrect return statements: 'return' and 'yield' outside functions
  */
 public class ReturnAnnotator extends PyAnnotator {
   public void visitPyReturnStatement(final PyReturnStatement node) {
@@ -36,8 +36,13 @@ public class ReturnAnnotator extends PyAnnotator {
     if (!(owner instanceof PyFunction || owner instanceof PyLambdaExpression)) {
       getHolder().createErrorAnnotation(node, "'yield' outside of function");
     }
-    if (owner instanceof PyFunction && ((PyFunction)owner).isAsync()) {
-      getHolder().createErrorAnnotation(node, "'yield' inside async function");
+
+    if (node.isDelegating() && owner instanceof PyFunction) {
+      final PyFunction function = (PyFunction)owner;
+
+      if (function.isAsync() && function.isAsyncAllowed()) {
+        getHolder().createErrorAnnotation(node, "Python does not support 'yield from' inside async functions");
+      }
     }
   }
 }

python/testData/highlighting/yieldInsideAsyncDef.py

@@ -1,5 +0,0 @@
-async def foo(x):
-    await x
-    <error descr="'yield' inside async function">yield x</error>
-    <error descr="'yield' inside async function">yield from x</error>
-    return x

python/testData/highlighting/yieldInsideAsyncDefPy35.py

@@ -0,0 +1,5 @@
+async def foo(x):
+    await x
+    <error descr="Python version 3.5 does not support 'yield' inside async functions">yield x</error>
+    <error descr="Python does not support 'yield from' inside async functions">yield from x</error>
+    return x

python/testData/highlighting/yieldInsideAsyncDefPy36.py

@@ -0,0 +1,5 @@
+async def foo(x):
+    await x
+    yield x
+    <error descr="Python does not support 'yield from' inside async functions">yield from x</error>
+    return x

python/testData/inspections/PyCompatibilityInspection/yieldInsideAsyncDef.py

@@ -0,0 +1,5 @@
+<warning descr="Python versions < 3.5 do not support this syntax">async</warning> def foo(x):
+    <warning descr="Python versions < 3.5 do not support this syntax">await x</warning>
+    <warning descr="Python version 3.5 does not support 'yield' inside async functions">yield x</warning>
+    <error descr="Python does not support 'yield from' inside async functions"><warning descr="Python versions < 3.3 do not support this syntax. Delegating to a subgenerator is available since Python 3.3; use explicit iteration over subgenerator instead.">yield from x</warning></error>
+    <warning descr="Python versions < 3.3 do not allow 'return' with argument inside generator.">return x</warning>

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

@@ -258,10 +258,15 @@ public class PythonHighlightingTest extends PyTestCase {
     doTest(LanguageLevel.PYTHON35, true, false);
   }
 
-  public void testYieldInsideAsyncDef() {
+  public void testYieldInsideAsyncDefPy35() {
     doTest(LanguageLevel.PYTHON35, false, false);
   }
 
+  // PY-20770
+  public void testYieldInsideAsyncDefPy36() {
+    doTest(LanguageLevel.PYTHON36, true, false);
+  }
+
   public void testUnpackingStar() {
     doTest(LanguageLevel.PYTHON35, false, false);
   }

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

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