contextlib.nested inconsistent with, well, nested with statements due exceptions raised in __enter__

Python/compile.c

Index: Python/compile.c
===================================================================
--- Python/compile.c	(revision 70381)
+++ Python/compile.c	(working copy)
@@ -2802,127 +2802,233 @@
 }
 
 /*
-   Implements the with statement from PEP 343.
+   Implements the with statement from PEP 343 and PEP 377.
 
-   The semantics outlined in that PEP are as follows:  
+   The semantics outlined in those PEPs are as follows:
 
-   with EXPR as VAR:
-       BLOCK
-  
+    with EXPR as VAR:
+        BLOCK
+
    It is implemented roughly as:
-  
-   context = EXPR
-   exit = context.__exit__  # not calling it
-   value = context.__enter__()
-   try:
-       VAR = value  # if VAR present in the syntax
-       BLOCK
-   finally:
-       if an exception was raised:
-	   exc = copy of (exception, instance, traceback)
-       else:
-	   exc = (None, None, None)
-       exit(*exc)
+
+    mgr = (EXPR)
+    exit = mgr.__exit__  # Not calling it yet
+    try:
+        value = mgr.__enter__()
+    except SkipStatement:
+        VAR = StatementSkipped
+        # Only if "as VAR" is present and
+        # VAR is a single name
+        # If VAR is a tuple of names, then StatementSkipped
+        # will be assigned to each name in the tuple
+    else:
+        call_exit = True
+        try:
+            try:
+                VAR = value  # Only if "as VAR" is present
+                BLOCK
+            except:
+                # The exceptional case is handled here
+                call_exit = False
+                if not exit(*sys.exc_info()):
+                    raise
+                # The exception is swallowed if exit() returns true
+        finally:
+            # The normal and non-local-goto cases are handled here
+            if call_exit:
+                exit(None, None, None)
  */
 static int
 compiler_with(struct compiler *c, stmt_ty s)
 {
-    static identifier enter_attr, exit_attr;
-    basicblock *block, *finally;
-    identifier tmpvalue = NULL;
+    static identifier enter_attr, exit_attr, skip_exc, skipped;
+    basicblock *outer_try, *outer_except, *outer_reraise;
+    basicblock *outer_else, *outer_end;
+    basicblock *inner_try, *inner_finally;
+    identifier tmpexit, tmpenter, tmpvalue = NULL;
+    expr_ty vars_expr;
 
     assert(s->kind == With_kind);
 
     if (!enter_attr) {
-	enter_attr = PyString_InternFromString("__enter__");
-	if (!enter_attr)
-	    return 0;
+        enter_attr = PyString_InternFromString("__enter__");
+        if (!enter_attr)
+            return 0;
     }
     if (!exit_attr) {
-	exit_attr = PyString_InternFromString("__exit__");
-	if (!exit_attr)
-	    return 0;
+        exit_attr = PyString_InternFromString("__exit__");
+        if (!exit_attr)
+            return 0;
     }
+    if (!skip_exc) {
+        skip_exc = PyString_InternFromString("SkipStatement");
+        if (!skip_exc)
+            return 0;
+    }
 
-    block = compiler_new_block(c);
-    finally = compiler_new_block(c);
-    if (!block || !finally)
-	return 0;
+    if (!skipped) {
+        skipped = PyString_InternFromString("StatementSkipped");
+        if (!skipped)
+            return 0;
+    }
 
-    if (s->v.With.optional_vars) {
-	/* Create a temporary variable to hold context.__enter__().
-	   We need to do this rather than preserving it on the stack
-	   because SETUP_FINALLY remembers the stack level.
-	   We need to do the assignment *inside* the try/finally
-	   so that context.__exit__() is called when the assignment
-	   fails.  But we need to call context.__enter__() *before*
-	   the try/finally so that if it fails we won't call
-	   context.__exit__().
-	*/
-	tmpvalue = compiler_new_tmpname(c);
-	if (tmpvalue == NULL)
-	    return 0;
-	PyArena_AddPyObject(c->c_arena, tmpvalue);
+
+    outer_try = compiler_new_block(c);
+    outer_except = compiler_new_block(c);
+    outer_reraise = compiler_new_block(c);
+    outer_else = compiler_new_block(c);
+    inner_try = compiler_new_block(c);
+    inner_finally = compiler_new_block(c);
+    outer_end = compiler_new_block(c);
+    if (!outer_try || !outer_except || !outer_reraise || !outer_else ||
+        !inner_try || !inner_finally || !outer_end)
+        return 0;
+
+    /* Create temporary variables to hold the context.__exit__ and
+     context.__enter__ methods and the context.__enter__() result.
+     These can't be stored on the stack because SETUP_EXCEPT and
+     SETUP_FINALLY remember the stack level and we want to retrieve
+     these values outside the relevant try blocks, but use them
+     *inside* those blocks. If we try to keep them on the stack
+     then the stack accounting gets out of whack when an
+     exception occurs. Thus, the use of temporary variables.
+     */
+    tmpexit = compiler_new_tmpname(c);
+    if (tmpexit == NULL)
+        return 0;
+    PyArena_AddPyObject(c->c_arena, tmpexit);
+    tmpenter = compiler_new_tmpname(c);
+    if (tmpenter == NULL)
+        return 0;
+    PyArena_AddPyObject(c->c_arena, tmpenter);
+    vars_expr = s->v.With.optional_vars;
+    if (vars_expr) {
+        tmpvalue = compiler_new_tmpname(c);
+        if (tmpvalue == NULL)
+            return 0;
+        PyArena_AddPyObject(c->c_arena, tmpvalue);
     }
 
     /* Evaluate EXPR */
     VISIT(c, expr, s->v.With.context_expr);
 
-    /* Squirrel away context.__exit__ by stuffing it under context */
+    /* XXX (ncoghlan): The with statement setup code has reached a
+     point were it is just crying out for a SETUP_WITH opcode.
+     Everything from here to the "BLOCK code" comment is bytecode
+     generated entirely by the compiler - that is a *lot* of trips
+     around the eval loop which could conceivably be avoided
+     Some of the temporary variables could go away as well.
+    */
+    /* Squirrel away context.__exit__ */
     ADDOP(c, DUP_TOP);
     ADDOP_O(c, LOAD_ATTR, exit_attr, names);
-    ADDOP(c, ROT_TWO);
+    if (!compiler_nameop(c, tmpexit, Store))
+        return 0;
 
+    /* Squirrel away context.__enter__ */
+    ADDOP_O(c, LOAD_ATTR, enter_attr, names);
+    if (!compiler_nameop(c, tmpenter, Store))
+        return 0;
+
+    /* Set up the outer try/except that allows context.__enter__ to
+     skip the body of the with statement and resume execution
+     with the following statement
+    */
+    ADDOP_JREL(c, SETUP_EXCEPT, outer_except);
+    compiler_use_next_block(c, outer_try);
+    if (!compiler_push_fblock(c, EXCEPT, outer_try))
+        return 0;
+
     /* Call context.__enter__() */
-    ADDOP_O(c, LOAD_ATTR, enter_attr, names);
+    if (!compiler_nameop(c, tmpenter, Load) ||
+        !compiler_nameop(c, tmpenter, Del))
+        return 0;
     ADDOP_I(c, CALL_FUNCTION, 0);
 
-    if (s->v.With.optional_vars) {
-	/* Store it in tmpvalue */
-	if (!compiler_nameop(c, tmpvalue, Store))
-	    return 0;
+    if (vars_expr) {
+        /* Store result in tmpvalue */
+        if (!compiler_nameop(c, tmpvalue, Store))
+            return 0;
+    } else {
+        /* Discard result from context.__enter__() */
+        ADDOP(c, POP_TOP);
     }
-    else {
-	/* Discard result from context.__enter__() */
-	ADDOP(c, POP_TOP);
+
+    ADDOP(c, POP_BLOCK);
+    compiler_pop_fblock(c, EXCEPT, outer_try);
+    ADDOP_JREL(c, JUMP_FORWARD, outer_else);
+
+    /* Specific exception handler for SkipStatement */
+    compiler_use_next_block(c, outer_except);
+    ADDOP(c, DUP_TOP);
+    ADDOP_O(c, LOAD_GLOBAL, skip_exc, names);
+    ADDOP_I(c, COMPARE_OP, PyCmp_EXC_MATCH);
+    ADDOP_JABS(c, POP_JUMP_IF_FALSE, outer_reraise);
+    ADDOP(c, POP_TOP);
+    ADDOP(c, POP_TOP);
+    ADDOP(c, POP_TOP);
+    if (vars_expr) {
+        /* Assign StatementSkipped to all VAR targets */
+        ADDOP_O(c, LOAD_GLOBAL, skipped, names);
+        if (vars_expr->kind == Tuple_kind) {
+            int i;
+            int n = asdl_seq_LEN(vars_expr->v.Tuple.elts);
+            for (i = 1; i < n; i++) {
+                ADDOP(c, DUP_TOP);
+            }
+            ADDOP_I(c, BUILD_TUPLE, n);
+        }
+        VISIT(c, expr, s->v.With.optional_vars);
     }
+    ADDOP_JREL(c, JUMP_FORWARD, outer_end);
 
-    /* Start the try block */
-    ADDOP_JREL(c, SETUP_FINALLY, finally);
+    /* Jump target for exceptions that don't match SkipStatement */
+    compiler_use_next_block(c, outer_reraise);
+    ADDOP(c, END_FINALLY);
 
-    compiler_use_next_block(c, block);
-    if (!compiler_push_fblock(c, FINALLY_TRY, block)) {
-	return 0;
+    /* Start the outer else clause */
+    compiler_use_next_block(c, outer_else);
+
+    /* Start the inner try block */
+    ADDOP_JREL(c, SETUP_FINALLY, inner_finally);
+
+    compiler_use_next_block(c, inner_try);
+    if (!compiler_push_fblock(c, FINALLY_TRY, inner_try)) {
+        return 0;
     }
 
     if (s->v.With.optional_vars) {
-	/* Bind saved result of context.__enter__() to VAR */
-	if (!compiler_nameop(c, tmpvalue, Load) ||
-	    !compiler_nameop(c, tmpvalue, Del))
-	  return 0;
-	VISIT(c, expr, s->v.With.optional_vars);
+        /* Bind saved result of context.__enter__() to VAR */
+        if (!compiler_nameop(c, tmpvalue, Load) ||
+            !compiler_nameop(c, tmpvalue, Del))
+            return 0;
+        VISIT(c, expr, s->v.With.optional_vars);
     }
 
     /* BLOCK code */
     VISIT_SEQ(c, stmt, s->v.With.body);
 
-    /* End of try block; start the finally block */
+    /* End of inner try block; start the finally block */
     ADDOP(c, POP_BLOCK);
-    compiler_pop_fblock(c, FINALLY_TRY, block);
+    compiler_pop_fblock(c, FINALLY_TRY, inner_try);
 
     ADDOP_O(c, LOAD_CONST, Py_None, consts);
-    compiler_use_next_block(c, finally);
-    if (!compiler_push_fblock(c, FINALLY_END, finally))
-	return 0;
+    compiler_use_next_block(c, inner_finally);
+    if (!compiler_push_fblock(c, FINALLY_END, inner_finally))
+        return 0;
 
-    /* Finally block starts; context.__exit__ is on the stack under
-       the exception or return information. Just issue our magic
-       opcode. */
+    /* Finally block starts; push tmpexit and issue our magic opcode. */
+    if (!compiler_nameop(c, tmpexit, Load) ||
+        !compiler_nameop(c, tmpexit, Del))
+        return 0;
     ADDOP(c, WITH_CLEANUP);
 
     /* Finally block ends. */
     ADDOP(c, END_FINALLY);
-    compiler_pop_fblock(c, FINALLY_END, finally);
+    compiler_pop_fblock(c, FINALLY_END, inner_finally);
+
+    /* Jump target for the end of the try/except statement */
+    compiler_use_next_block(c, outer_end);
     return 1;
 }