Py CF improvements: handle except parts more accurately

python/src/com/jetbrains/python/codeInsight/controlflow/PyControlFlowBuilder.java

@@ -409,6 +409,7 @@ public class PyControlFlowBuilder extends PyRecursiveElementVisitor {
     for (Pair<PsiElement, Instruction> pair : myPending) {
       myBuilder.addPendingEdge(pair.first, pair.second);
     }
+    // Finally part handling
     final PyFinallyPart finallyPart = node.getFinallyPart();
     Instruction finallyInstruction = null;
     Instruction lastFinallyInstruction = null;
@@ -424,30 +425,40 @@ public class PyControlFlowBuilder extends PyRecursiveElementVisitor {
     myBuilder.processPending(new ControlFlowBuilder.PendingProcessor() {
       public void process(final PsiElement pendingScope, final Instruction instruction) {
         final PsiElement pendingElement = instruction.getElement();
+        if (pendingElement == null){
+          myBuilder.addPendingEdge(pendingScope, instruction);
+          return;
+        }
 
-        // handle raise instructions inside compound statement
-        if (pendingElement instanceof PyRaiseStatement && PsiTreeUtil.isAncestor(tryPart, pendingElement, false)) {
+        // Process raise statements
+        if (PsiTreeUtil.isAncestor(tryPart, pendingElement, false) &&
+            PsiTreeUtil.getParentOfType(pendingElement, PyRaiseStatement.class) != null){
           for (Instruction rescueInstruction : exceptInstructions) {
             myBuilder.addEdge(instruction, rescueInstruction);
           }
-          return;
-        }
-        // handle return pending instructions inside body if ensure block exists
-        if (pendingElement instanceof PyReturnStatement && !finallyRef.isNull() && PsiTreeUtil.isAncestor(node, pendingElement, false)) {
-          myBuilder.addEdge(instruction, finallyRef.get());
-          myBuilder.addPendingEdge(null, lastFinallyRef.get());
+          myBuilder.addPendingEdge(pendingScope, instruction);
           return;
         }
 
-        // Handle pending instructions inside body with ensure block
-        if (pendingElement != null &&
-            finallyPart != null &&
-            pendingScope != finallyPart &&
-            PsiTreeUtil.isAncestor(node, pendingElement, false)) {
+        // Add except statements
+        if (PsiTreeUtil.isAncestor(tryPart, pendingElement, false)){
+          for (Instruction rescueInstruction : exceptInstructions) {
+            myBuilder.addEdge(instruction, rescueInstruction);
+          }
+        }
+
+        // Process finally
+        if (!finallyRef.isNull()){
           myBuilder.addEdge(instruction, finallyRef.get());
+          if (!lastFinallyRef.isNull() &&
+              PsiTreeUtil.getParentOfType(pendingElement, PyReturnStatement.class, false) != null){
+            myBuilder.addPendingEdge(null, lastFinallyRef.get());
+          }
           return;
         }
+
         myBuilder.addPendingEdge(pendingScope, instruction);
+
       }
     });
   }

python/src/com/jetbrains/python/inspections/PyUnboundLocalVariableInspection.java

@@ -52,11 +52,6 @@ public class PyUnboundLocalVariableInspection extends LocalInspectionTool {
     return new PyInspectionVisitor(holder){
       @Override
       public void visitPyReferenceExpression(final PyReferenceExpression node) {
-        if (PsiTreeUtil.getParentOfType(node, PyExceptPart.class) != null){
-          // TODO[oleg] more accurate check
-          return;
-        }
-
         // Ignore callee expressions
         if (PyCallExpressionNavigator.getPyCallExpressionByCallee(node) != null){
           return;

python/testData/codeInsight/controlflow/fortrycontinue.py

@@ -0,0 +1,7 @@
+tests = "foo"
+for t in tests:
+    try:
+        for t in []:
+            print t
+    except Exception:
+        continue

python/testData/codeInsight/controlflow/fortrycontinue.txt

@@ -0,0 +1,18 @@
+0(1) element: null
+1(2) element: PyAssignmentStatement
+2(3) WRITE ACCESS: tests
+3(4) element: PyForStatement
+4(5,17) READ ACCESS: tests
+5(6) element: PyStatementList
+6(7) WRITE ACCESS: t
+7(8) element: PyTryExceptStatement
+8(9) element: PyTryPart
+9(10,14,5,17) element: PyForStatement
+10(11) element: PyStatementList
+11(12) WRITE ACCESS: t
+12(13) element: PyPrintStatement
+13(10,14,5,17) READ ACCESS: t
+14(15) element: PyExceptPart
+15(16) READ ACCESS: Exception
+16(3) element: PyContinueStatement
+17() element: null

python/testData/codeInsight/controlflow/ifelsereturn.txt

@@ -5,7 +5,7 @@
 4(5) element: PyTryExceptStatement
 5(6) element: PyTryPart
 6(7) element: PyReturnStatement
-7(15) READ ACCESS: renderer
+7(8,15) READ ACCESS: renderer
 8(9) element: PyExceptPart
 9(10) READ ACCESS: KeyError
 10(11) element: PyRaiseStatement

python/testData/codeInsight/controlflow/raise.txt

@@ -2,7 +2,7 @@
 1(2) element: PyTryExceptStatement
 2(3) element: PyTryPart
 3(4) element: PyRaiseStatement
-4(11) READ ACCESS: Exception
+4(5,11) READ ACCESS: Exception
 5(6) element: PyExceptPart
 6(7) element: PyAssignmentStatement
 7(8) READ ACCESS: Exception

python/testData/codeInsight/controlflow/trybreak.txt

@@ -2,10 +2,10 @@
 1(2) element: PyTryExceptStatement
 2(3) element: PyTryPart
 3(4) element: PyForStatement
-4(5,11) READ ACCESS: bar
+4(5,8,11) READ ACCESS: bar
 5(6) element: PyStatementList
 6(7) WRITE ACCESS: i
-7(11) element: PyBreakStatement
+7(8,11) element: PyBreakStatement
 8(9) element: PyExceptPart
 9(10) element: PyRaiseStatement
 10(12) READ ACCESS: Exception

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

@@ -98,6 +98,10 @@ public class PyControlFlowBuilderTest extends LightMarkedTestCase {
     doTest();
   }
 
+  public void testForTryContinue() throws Exception {
+    doTest();
+  }
+
   public void testTryBreak() throws Exception {
     final String testName = getTestName(false).toLowerCase();
     configureByFile(testName + ".py");