Yet another piece of Python control flow

python/src/com/jetbrains/python/psi/PyBreakStatement.java

@@ -16,6 +16,8 @@
 
 package com.jetbrains.python.psi;
 
+import org.jetbrains.annotations.Nullable;
+
 /**
  * Created by IntelliJ IDEA.
  * User: yole
@@ -24,4 +26,6 @@ package com.jetbrains.python.psi;
  * To change this template use File | Settings | File Templates.
  */
 public interface PyBreakStatement extends PyStatement {
+  @Nullable
+  PyLoopStatement getLoopStatement();
 }

python/src/com/jetbrains/python/psi/PyContinueStatement.java

@@ -16,6 +16,8 @@
 
 package com.jetbrains.python.psi;
 
+import org.jetbrains.annotations.Nullable;
+
 /**
  * Created by IntelliJ IDEA.
  * User: yole
@@ -24,4 +26,6 @@ package com.jetbrains.python.psi;
  * To change this template use File | Settings | File Templates.
  */
 public interface PyContinueStatement extends PyStatement {
+  @Nullable
+  PyLoopStatement getLoop();
 }

python/src/com/jetbrains/python/psi/impl/PyBreakStatementImpl.java

@@ -17,8 +17,11 @@
 package com.jetbrains.python.psi.impl;
 
 import com.intellij.lang.ASTNode;
+import com.intellij.psi.util.PsiTreeUtil;
 import com.jetbrains.python.psi.PyBreakStatement;
 import com.jetbrains.python.psi.PyElementVisitor;
+import com.jetbrains.python.psi.PyLoopStatement;
+import org.jetbrains.annotations.Nullable;
 
 /**
  * Created by IntelliJ IDEA.
@@ -28,11 +31,17 @@ import com.jetbrains.python.psi.PyElementVisitor;
  * To change this template use File | Settings | File Templates.
  */
 public class PyBreakStatementImpl extends PyElementImpl implements PyBreakStatement {
-    public PyBreakStatementImpl(ASTNode astNode) {
-        super(astNode);
-    }
+  public PyBreakStatementImpl(ASTNode astNode) {
+    super(astNode);
+  }
 
-    @Override protected void acceptPyVisitor(PyElementVisitor pyVisitor) {
-        pyVisitor.visitPyBreakStatement(this);
-    }
+  @Override
+  protected void acceptPyVisitor(PyElementVisitor pyVisitor) {
+    pyVisitor.visitPyBreakStatement(this);
+  }
+
+  @Nullable
+  public PyLoopStatement getLoopStatement() {
+    return PsiTreeUtil.getParentOfType(this, PyLoopStatement.class);
+  }
 }

python/src/com/jetbrains/python/psi/impl/PyContinueStatementImpl.java

@@ -17,8 +17,11 @@
 package com.jetbrains.python.psi.impl;
 
 import com.intellij.lang.ASTNode;
+import com.intellij.psi.util.PsiTreeUtil;
 import com.jetbrains.python.psi.PyContinueStatement;
 import com.jetbrains.python.psi.PyElementVisitor;
+import com.jetbrains.python.psi.PyLoopStatement;
+import org.jetbrains.annotations.Nullable;
 
 /**
  * Created by IntelliJ IDEA.
@@ -35,4 +38,9 @@ public class PyContinueStatementImpl extends PyElementImpl implements PyContinue
     @Override protected void acceptPyVisitor(PyElementVisitor pyVisitor) {
         pyVisitor.visitPyContinueStatement(this);
     }
+
+  @Nullable
+  public PyLoopStatement getLoop() {
+    return PsiTreeUtil.getParentOfType(this, PyLoopStatement.class);
+  }
 }

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

@@ -20,23 +20,34 @@ public class PyControlFlowBuilderTest extends LightMarkedTestCase {
     return PythonTestUtil.getTestDataPath() + "/psi/controlflow/";
   }
 
-
-  public void testFile() throws Exception {
-    check("file.py", "file.txt");
-  }
-
-  private void check(final String inputFile, final String outputFile) throws Exception {
-    configureByFile(inputFile);
+  private void doTest() throws Exception {
+    final String testName = getTestName(false).toLowerCase();
+    configureByFile(testName + ".py");
     final StringBuffer buffer = new StringBuffer();
-    final ControlFlow flow = ((PyFile)myFile).getControlFlow();    
+    final ControlFlow flow = ((PyFile)myFile).getControlFlow();
     final Instruction[] instructions = flow.getInstructions();
     for (Instruction instruction : instructions) {
       buffer.append(instruction).append("\n");
     }
-    final String fullPath = getTestDataPath() + outputFile;
+    final String fullPath = getTestDataPath() + testName + ".txt";
     final VirtualFile vFile = LocalFileSystem.getInstance().findFileByPath(fullPath.replace(File.separatorChar, '/'));
     final String fileText = StringUtil.convertLineSeparators(VfsUtil.loadText(vFile), "\n");
     assertEquals(fileText.trim(), buffer.toString().trim());
   }
+  
+  public void testFile() throws Exception {
+    doTest();
+  }
+  
+  public void testIf() throws Exception {
+    doTest();
+  }
 
+  public void testFor() throws Exception {
+    doTest();
+  }
+
+  public void testWhile() throws Exception {
+    doTest();
+  }
 }