PY-992 Cannot perform refactoring when execution flow is interrupted

python/src/com/jetbrains/python/codeInsight/codeFragment/PyCodeFragmentUtil.java

@@ -92,23 +92,9 @@ public class PyCodeFragmentUtil {
     }
 
     // If we see more than 1 outer instruction, controlflow is interrupted
-    if (outerInstructions.size() > 2){
+    if (outerInstructions.size() > 1){
       throw new CannotCreateCodeFragmentException(PyBundle.message("refactoring.extract.method.error.cannot.perform.refactoring.when.execution.flow.is.interrupted"));
     }
-    if (outerInstructions.size() == 2){
-      boolean errorFound = true;
-      for (Instruction outerInstruction : outerInstructions) {
-        // Here we check control flow when for statement content is beeing extracted
-        if (outerInstruction instanceof ReadWriteInstruction &&
-            PyForStatementNavigator.getPyForStatementByIterable(outerInstruction.getElement())!=null){
-          errorFound = false;
-          break;
-        }
-      }
-      if (errorFound){
-        throw new CannotCreateCodeFragmentException(PyBundle.message("refactoring.extract.method.error.cannot.perform.refactoring.when.execution.flow.is.interrupted"));
-      }
-    }
 
     // Building code fragment
     final PyCodeFragmentBuilder builder = new PyCodeFragmentBuilder(owner, start, end);

python/testData/codeInsight/codefragment/forifreturn.test

@@ -0,0 +1,4 @@
+for start, end, name in self.talks:
+    <begin>if hour >= start and hour < end: return name<end>
+<result>
+Cannot perform refactoring when execution flow is interrupted

python/testData/codeInsight/controlflow/forreturn.py

@@ -0,0 +1,2 @@
+for start, end, name in self.talks:
+  if hour >= start and hour < end: return name

python/testData/codeInsight/controlflow/forreturn.txt

@@ -0,0 +1,16 @@
+0(1) element: null
+1(2) element: PyForStatement
+2(3,15) READ ACCESS: self
+3(4) element: PyStatementList
+4(5) WRITE ACCESS: start
+5(6) WRITE ACCESS: end
+6(7) WRITE ACCESS: name
+7(8) element: PyIfStatement
+8(9) READ ACCESS: hour
+9(10) READ ACCESS: start
+10(11) READ ACCESS: hour
+11(12,3,15) READ ACCESS: end
+12(13) element: PyStatementList. Condition: hour >= start and hour < end:true
+13(14) element: PyReturnStatement
+14(15) READ ACCESS: name
+15() element: null

python/testSrc/com/jetbrains/python/PyControlFlowBuilderTest.java

@@ -103,6 +103,10 @@ public class PyControlFlowBuilderTest extends LightMarkedTestCase {
     doTest();
   }
 
+  public void testForReturn() throws Exception {
+    doTest();
+  }
+
   public void testForTryContinue() throws Exception {
     doTest();
   }

python/testSrc/com/jetbrains/python/refactoring/PyCodeFragmentTest.java

@@ -1,5 +1,6 @@
 package com.jetbrains.python.refactoring;
 
+import com.intellij.codeInsight.codeFragment.CannotCreateCodeFragmentException;
 import com.intellij.codeInsight.codeFragment.CodeFragment;
 import com.intellij.openapi.util.Pair;
 import com.intellij.openapi.util.text.StringUtil;
@@ -67,19 +68,24 @@ public class PyCodeFragmentTest extends LightMarkedTestCase {
       context = PsiTreeUtil.getParentOfType(context, ScopeOwner.class);
     }
     final StringBuffer buffer = new StringBuffer();
-    final CodeFragment fragment = PyCodeFragmentUtil.createCodeFragment((ScopeOwner)context, startElement, endElement);
-    if (fragment.isReturnInstructonInside()) {
-      buffer.append("Return instruction inside found").append("\n");
+    try {
+      final CodeFragment fragment = PyCodeFragmentUtil.createCodeFragment((ScopeOwner)context, startElement, endElement);
+      if (fragment.isReturnInstructonInside()) {
+        buffer.append("Return instruction inside found").append("\n");
+      }
+      buffer.append("In:\n");
+      for (String inputVariable : new TreeSet<String>(fragment.getInputVariables())) {
+        buffer.append(inputVariable).append('\n');
+      }
+      buffer.append("Out:\n");
+      for (String outputVariable : new TreeSet<String>(fragment.getOutputVariables())) {
+        buffer.append(outputVariable).append('\n');
+      }
     }
-    buffer.append("In:\n");
-    for (String inputVariable : new TreeSet<String>(fragment.getInputVariables())) {
-      buffer.append(inputVariable).append('\n');
+    catch (CannotCreateCodeFragmentException e) {
+      assertEquals(result.trim(), e.getMessage());
+      return;
     }
-    buffer.append("Out:\n");
-    for (String outputVariable : new TreeSet<String>(fragment.getOutputVariables())) {
-      buffer.append(outputVariable).append('\n');
-    }
-
     assertEquals(result.trim(), buffer.toString().trim());
   }
 
@@ -144,5 +150,9 @@ public class PyCodeFragmentTest extends LightMarkedTestCase {
   public void testClass() throws Exception {
     doTest();
   }
+
+  public void testForIfReturn() throws Exception {
+    doTest();
+  }
 }