Added intention for replacing explicit iteration with delegation to sub-generator for Python 3.3 (PY-7383)

python/src/META-INF/python-plugin-common.xml

@@ -240,6 +240,11 @@
       <category>Python</category>
     </intentionAction>
 
+    <intentionAction>
+      <className>com.jetbrains.python.codeInsight.intentions.PyYieldFromIntention</className>
+      <category>Python</category>
+    </intentionAction>
+
     <testFinder implementation="com.jetbrains.python.codeInsight.testIntegration.PyTestFinder"/>
     <testCreator language="Python" implementationClass="com.jetbrains.python.codeInsight.testIntegration.PyTestCreator"/>
 

python/src/com/jetbrains/python/PyBundle.properties

@@ -191,6 +191,9 @@ INTN.specify.type.in.annotation=Specify type for reference using annotation
 #TypeAssertionIntention
 INTN.insert.assertion=Insert type assertion
 
+#PyYieldFromIntention
+INTN.yield.from=Transform explicit iteration with 'yield' into 'yield from' expression
+
 # Conflict checker
 CONFLICT.name.$0.obscured=Name ''{0}'' obscured by local definitions
 CONFLICT.name.$0.obscured.cannot.convert=Name ''{0}'' obscured. Cannot convert.

python/src/com/jetbrains/python/codeInsight/intentions/PyYieldFromIntention.java

@@ -0,0 +1,108 @@
+package com.jetbrains.python.codeInsight.intentions;
+
+import com.intellij.codeInsight.intention.impl.BaseIntentionAction;
+import com.intellij.openapi.editor.Editor;
+import com.intellij.openapi.project.Project;
+import com.intellij.psi.PsiElement;
+import com.intellij.psi.PsiFile;
+import com.intellij.psi.util.PsiTreeUtil;
+import com.intellij.util.IncorrectOperationException;
+import com.jetbrains.python.PyBundle;
+import com.jetbrains.python.psi.*;
+import org.jetbrains.annotations.NotNull;
+import org.jetbrains.annotations.Nullable;
+
+/**
+ * @author vlan
+ */
+public class PyYieldFromIntention extends BaseIntentionAction {
+  @NotNull
+  @Override
+  public String getFamilyName() {
+    return PyBundle.message("INTN.yield.from");
+  }
+
+  @NotNull
+  @Override
+  public String getText() {
+    return PyBundle.message("INTN.yield.from");
+  }
+
+  @Override
+  public boolean isAvailable(@NotNull Project project, Editor editor, PsiFile file) {
+    if (LanguageLevel.forElement(file).isAtLeast(LanguageLevel.PYTHON33)) {
+      final PyForStatement forLoop = findForStatementAtCaret(editor, file);
+      if (forLoop != null) {
+        final PyTargetExpression forTarget = findSingleForLoopTarget(forLoop);
+        final PyReferenceExpression yieldValue = findSingleYieldValue(forLoop);
+        if (forTarget != null && yieldValue != null) {
+          final String targetName = forTarget.getName();
+          if (targetName != null && targetName.equals(yieldValue.getName())) {
+            return true;
+          }
+        }
+      }
+    }
+    return false;
+  }
+
+  @Override
+  public void invoke(@NotNull Project project, Editor editor, PsiFile file) throws IncorrectOperationException {
+    final PyForStatement forLoop = findForStatementAtCaret(editor, file);
+    if (forLoop != null) {
+      final PyExpression source = forLoop.getForPart().getSource();
+      if (source != null) {
+        final PyElementGenerator generator = PyElementGenerator.getInstance(project);
+        final String text = "yield from foo";
+        final PyExpressionStatement exprStmt = generator.createFromText(LanguageLevel.forElement(file), PyExpressionStatement.class, text);
+        final PyExpression expr = exprStmt.getExpression();
+        if (expr instanceof PyYieldExpression) {
+          final PyExpression yieldValue = ((PyYieldExpression)expr).getExpression();
+          if (yieldValue != null) {
+            yieldValue.replace(source);
+            forLoop.replace(exprStmt);
+          }
+        }
+      }
+    }
+  }
+
+  @Nullable
+  private static PyForStatement findForStatementAtCaret(@NotNull Editor editor, @NotNull PsiFile file) {
+    final PsiElement elementAtCaret = file.findElementAt(editor.getCaretModel().getOffset());
+    return PsiTreeUtil.getParentOfType(elementAtCaret, PyForStatement.class);
+  }
+
+  @Nullable
+  private static PyTargetExpression findSingleForLoopTarget(@NotNull PyForStatement forLoop) {
+    final PyForPart forPart = forLoop.getForPart();
+    final PyExpression forTarget = forPart.getTarget();
+    if (forTarget instanceof PyTargetExpression) {
+      return (PyTargetExpression)forTarget;
+    }
+    return null;
+  }
+
+  @Nullable
+  private static PyReferenceExpression findSingleYieldValue(@NotNull PyForStatement forLoop) {
+    final PyForPart forPart = forLoop.getForPart();
+    final PyStatementList stmtList = forPart.getStatementList();
+    if (stmtList != null && forLoop.getElsePart() == null) {
+      final PyStatement[] statements = stmtList.getStatements();
+      if (statements.length == 1) {
+        final PyStatement firstStmt = statements[0];
+        if (firstStmt instanceof PyExpressionStatement) {
+          final PyExpression firstExpr = ((PyExpressionStatement)firstStmt).getExpression();
+          if (firstExpr instanceof PyYieldExpression) {
+            final PyYieldExpression yieldExpr = (PyYieldExpression)firstExpr;
+            final PyExpression yieldValue = yieldExpr.getExpression();
+            if (yieldValue instanceof PyReferenceExpression) {
+              return (PyReferenceExpression)yieldValue;
+            }
+          }
+        }
+      }
+    }
+    return null;
+  }
+}

python/testData/intentions/afterYieldFrom.py

@@ -0,0 +1,4 @@
+def f(g):
+    yield 'begin'
+    yield from g()
+    print('end')

python/testData/intentions/beforeYieldFrom.py

@@ -0,0 +1,5 @@
+def f(g):
+    yield 'begin'
+    for x in g():
+        yield <caret>x
+    print('end')

python/testSrc/com/jetbrains/python/PyIntentionTest.java

@@ -261,6 +261,11 @@ public class PyIntentionTest extends PyTestCase {
     doDocStubTest();
   }
 
+  // PY-7383
+  public void testYieldFrom() {
+    doTest(PyBundle.message("INTN.yield.from"), LanguageLevel.PYTHON33);
+  }
+
   private void doDocStubTest() {
     CodeInsightSettings codeInsightSettings = CodeInsightSettings.getInstance();
     codeInsightSettings.JAVADOC_STUB_ON_ENTER = true;