Fix unbound local variable after exception inside with statement (PY-13919)

Change control flow to make it more correct for with statement. Handle special case for Unbound local variable inspection and add tests

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

@@ -868,9 +868,7 @@ public class PyControlFlowBuilder extends PyRecursiveElementVisitor {
           PsiTreeUtil.getParentOfType(element, PyRaiseStatement.class) != null) {
         myBuilder.addPendingEdge(node, instruction);
       }
-      else {
-        myBuilder.addPendingEdge(pendingScope, instruction);
-      }
+      myBuilder.addPendingEdge(pendingScope, instruction);
     });
   }
 

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

@@ -42,6 +42,7 @@ import com.jetbrains.python.psi.impl.PyGlobalStatementNavigator;
 import com.jetbrains.python.psi.resolve.PyResolveUtil;
 import org.jetbrains.annotations.Nls;
 import org.jetbrains.annotations.NotNull;
+import org.jetbrains.annotations.Nullable;
 
 import java.util.HashSet;
 import java.util.Set;
@@ -140,6 +141,9 @@ public class PyUnboundLocalVariableInspection extends PyInspection {
             return;
           }
         }
+        if (isWriteAndRaiseInsideWith(node, resolved)) {
+          return;
+        }
         if (PyUnreachableCodeInspection.hasAnyInterruptedControlFlowPaths(node)) {
           return;
         }
@@ -161,6 +165,27 @@ public class PyUnboundLocalVariableInspection extends PyInspection {
       }
     }
 
+    private static boolean isWriteAndRaiseInsideWith(@NotNull PyReferenceExpression node, @Nullable PsiElement resolvedElement) {
+      if (resolvedElement != null && !PyUtil.inSameFile(node, resolvedElement)) {
+        return false;
+      }
+      boolean isExceptionRaised = false;
+      boolean isUnderContextManager = false;
+      PsiElement firstWith = PsiTreeUtil.getParentOfType(resolvedElement, PyWithStatement.class, true, ScopeOwner.class);
+      if (firstWith != null && PsiTreeUtil.findChildOfType(firstWith, PyRaiseStatement.class) != null) {
+        isExceptionRaised = true;
+        PyWithStatement withStatement = (PyWithStatement)firstWith;
+        for (PyWithItem withItem : withStatement.getWithItems()) {
+          PyExpression contextManager = withItem.getExpression();
+          if (contextManager != null) {
+            isUnderContextManager = true;
+            break;
+          }
+        }
+      }
+      return isExceptionRaised && isUnderContextManager;
+    }
+
     private static boolean isFirstUnboundRead(@NotNull PyReferenceExpression node, @NotNull ScopeOwner owner) {
       final String nodeName = node.getReferencedName();
       final Scope scope = ControlFlowCache.getScope(owner);

python/testData/codeInsight/controlflow/withraiseexception.py

@@ -0,0 +1,7 @@
+
+with context_manager:
+    if c:
+        raise ValueError
+    val = c
+
+print(val)

python/testData/codeInsight/controlflow/withraiseexception.txt

@@ -0,0 +1,15 @@
+0(1) element: null
+1(2) element: PyWithStatement
+2(3) READ ACCESS: context_manager
+3(4) element: PyIfStatement
+4(5,9) READ ACCESS: c
+5(6) element: PyStatementList. Condition: c:true
+6(7) ASSERTTYPE ACCESS: c
+7(8) element: PyRaiseStatement
+8(12,14) READ ACCESS: ValueError
+9(10) element: PyAssignmentStatement
+10(11) READ ACCESS: c
+11(12) WRITE ACCESS: val
+12(13) element: PyPrintStatement
+13(14) READ ACCESS: val
+14() element: null

python/testData/inspections/PyUnboundLocalVariableInspection/RaiseInsideWith.py

@@ -0,0 +1,18 @@
+
+
+def warn(c):
+    with open('foo.txt', 'rb') as FILE:
+        if c:
+            raise ValueError
+        val = 1
+
+    return val  # raises UnboundLocalError
+
+
+def ok(context):
+    with context as c:
+        if c:
+            raise ValueError
+        val = 1
+
+    return val  # raises UnboundLocalError

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

@@ -279,6 +279,11 @@ public class PyControlFlowBuilderTest extends LightMarkedTestCase {
     doTest();
   }
 
+  // PY-13919
+  public void testWithRaiseException() {
+    doTest();
+  }
+
   private void doTestFirstStatement() {
     final String testName = getTestName(false).toLowerCase();
     configureByFile(testName + ".py");

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

@@ -177,6 +177,11 @@ public class PyUnboundLocalVariableInspectionTest extends PyInspectionTestCase {
     doTest();
   }
 
+  // PY-13919
+  public void testRaiseInsideWith() {
+    doTest();
+  }
+
   // PY-6114
   public void testUnboundUnreachable() {
     doTest();