# HG changeset patch
# User Kristján Valur Jónsson <kristjan@ccpgames.com>
# Date 1375884232 0
# Node ID 6443a5d0e7a89fe751e098338542d544f47d64a0
# Parent  89ce323357dbd80f8d08b9e8616cd65532ceff61
[mq]: 2013-08-07_12-03-37_r85060+.diff

diff --git a/Include/pyerrors.h b/Include/pyerrors.h
--- a/Include/pyerrors.h
+++ b/Include/pyerrors.h
@@ -144,6 +144,7 @@
 PyAPI_DATA(PyObject *) PyExc_GeneratorExit;
 PyAPI_DATA(PyObject *) PyExc_ArithmeticError;
 PyAPI_DATA(PyObject *) PyExc_LookupError;
+PyAPI_DATA(PyObject *) PyExc_ContextManagerExit;
 
 PyAPI_DATA(PyObject *) PyExc_AssertionError;
 PyAPI_DATA(PyObject *) PyExc_AttributeError;
diff --git a/Lib/contextlib.py b/Lib/contextlib.py
--- a/Lib/contextlib.py
+++ b/Lib/contextlib.py
@@ -4,7 +4,8 @@
 from collections import deque
 from functools import wraps
 
-__all__ = ["contextmanager", "closing", "ContextDecorator", "ExitStack", "ignored"]
+__all__ = ["contextmanager", "closing", "ContextDecorator", "ExitStack", "ignored",
+           "nested", "nested_delayed"]
 
 
 class ContextDecorator(object):
@@ -47,7 +48,8 @@
         try:
             return next(self.gen)
         except StopIteration:
-            raise RuntimeError("generator didn't yield")
+            # No yield.  Assume that the body isn't to be run.
+            raise ContextManagerExit
 
     def __exit__(self, type, value, traceback):
         if type is None:
@@ -116,6 +118,74 @@
     return helper
 
 
+@contextmanager
+def nested(*managers):
+    """Combine multiple context managers into a single nested context manager.
+
+    The one advantage of this function over the multiple manager form of the
+    with statement is that argument unpacking allows it to be
+    used with a variable number of context managers as follows:
+
+        with nested(*managers):
+            do_something()
+
+    """
+    def get_callable(mgr):
+        def callable():
+            return mgr
+        return callable
+    managers = [get_callable(m) for m in managers]
+    with nested_delayed(*managers) as vars:
+        yield vars
+
+@contextmanager
+def nested_delayed(*managers):
+    """Combine multiple context managers into a single nested context manager.
+    This version takes a sequence of callables, each callable being called
+    and returning a context manager on demand.  This ensures that inner
+    context managers are not created before outer ones have successfuly been
+    entered.
+
+    The advantage over using the built in nested context manager syntax is
+    that the list of context manager can be dynamically created.
+
+        with nested_delayed((lambda:open(file1)), (lambda:open(file2))) as f1, f2:
+            do_something(f1, f2)
+
+    """
+    exits = []
+    vars = []
+    exc = (None, None, None)
+    try:
+        for mgr_call in managers:
+            mgr = mgr_call()
+            exit = mgr.__exit__
+            enter = mgr.__enter__
+            vars.append(enter())
+            exits.append(exit)
+        yield vars
+    except:
+        exc = sys.exc_info()
+    finally:
+        while exits:
+            exit = exits.pop()
+            try:
+                if exit(*exc):
+                    exc = (None, None, None)
+            except:
+                exc = sys.exc_info()
+        if exc != (None, None, None):
+            # Don't rely on sys.exc_info() still containing
+            # the right information. Another exception may
+            # have been raised and caught by an exit method
+            raise exc[0](exc[1]).with_traceback(exc[2])
+        if len(vars) < len(managers):
+            # An error in an __enter__ method was handled by
+            # an outer __exit__ method.  In this case, the
+            # managed body is not executed.
+            raise ContextManagerExit
+
+
 class closing(object):
     """Context to automatically close something at the end of a block.
 
diff --git a/Lib/test/test_contextlib.py b/Lib/test/test_contextlib.py
--- a/Lib/test/test_contextlib.py
+++ b/Lib/test/test_contextlib.py
@@ -106,6 +106,228 @@
         baz = self._create_contextmanager_attribs()
         self.assertEqual(baz.__doc__, "Whee!")
 
+class NestedTestCase(unittest.TestCase):
+
+    # XXX This needs more work
+
+    def test_nested(self):
+        @contextmanager
+        def a():
+            yield 1
+        @contextmanager
+        def b():
+            yield 2
+        @contextmanager
+        def c():
+            yield 3
+        with nested(a(), b(), c()) as (x, y, z):
+            self.assertEqual(x, 1)
+            self.assertEqual(y, 2)
+            self.assertEqual(z, 3)
+
+    def test_nested_cleanup(self):
+        state = []
+        @contextmanager
+        def a():
+            state.append(1)
+            try:
+                yield 2
+            finally:
+                state.append(3)
+        @contextmanager
+        def b():
+            state.append(4)
+            try:
+                yield 5
+            finally:
+                state.append(6)
+        with self.assertRaises(ZeroDivisionError):
+            with nested(a(), b()) as (x, y):
+                state.append(x)
+                state.append(y)
+                1 // 0
+        self.assertEqual(state, [1, 4, 2, 5, 6, 3])
+
+    def test_nested_right_exception(self):
+        @contextmanager
+        def a():
+            yield 1
+        class b(object):
+            def __enter__(self):
+                return 2
+            def __exit__(self, *exc_info):
+                try:
+                    raise Exception()
+                except:
+                    pass
+        with self.assertRaises(ZeroDivisionError):
+            with nested(a(), b()) as (x, y):
+                1 // 0
+        self.assertEqual((x, y), (1, 2))
+
+    def test_nested_b_swallows(self):
+        @contextmanager
+        def a():
+            yield
+        @contextmanager
+        def b():
+            try:
+                yield
+            except:
+                # Swallow the exception
+                pass
+        try:
+            with nested(a(), b()):
+                1 // 0
+        except ZeroDivisionError:
+            self.fail("Didn't swallow ZeroDivisionError")
+
+    def test_nested_break(self):
+        @contextmanager
+        def a():
+            yield
+        state = 0
+        while True:
+            state += 1
+            with nested(a(), a()):
+                break
+            state += 10
+        self.assertEqual(state, 1)
+
+    def test_nested_continue(self):
+        @contextmanager
+        def a():
+            yield
+        state = 0
+        while state < 3:
+            state += 1
+            with nested(a(), a()):
+                continue
+            state += 10
+        self.assertEqual(state, 3)
+
+    def test_nested_return(self):
+        @contextmanager
+        def a():
+            try:
+                yield
+            except:
+                pass
+        def foo():
+            with nested(a(), a()):
+                return 1
+            return 10
+        self.assertEqual(foo(), 1)
+
+    def test_nested_handle_enter_error(self):
+        # Test that an inner __enter__() causing an error,
+        # which is then handled by an outer __exit__(),
+        # causes the correct behaviour, i.e. the body
+        # is not executed.
+        @contextmanager
+        def b():
+            # raise error when entering
+            raise ZeroDivisionError
+            yield
+        @contextmanager
+        def a():
+            # Handle the inner error
+            try:
+                yield
+            except ZeroDivisionError:
+                return True
+
+        def foo():
+            with nested(a(), b()):
+                return 1
+            return 10
+        self.assertEqual(foo(), 10)
+
+class NestedDelayedCase(unittest.TestCase):
+    """
+    Test the nested_delaye context manager which nests
+    context managers return from callables
+    """
+
+    def test_nested(self):
+        @contextmanager
+        def a():
+            yield 1
+        @contextmanager
+        def b():
+            yield 2
+        @contextmanager
+        def c():
+            yield 3
+        with nested_delayed(a, b, c) as (x, y, z):
+            self.assertEqual(x, 1)
+            self.assertEqual(y, 2)
+            self.assertEqual(z, 3)
+
+    def test_nested_fail(self):
+        # Test that an outer error causes the inner
+        # callable to not be called
+        c_ran = False
+        @contextmanager
+        def _c():
+            yield 1
+        def c():
+            nonlocal c_ran
+            c_ran = True
+            return _c()
+        @contextmanager
+        def b():
+            raise ZeroDivisionError
+            yield 2
+        @contextmanager
+        def a():
+            yield 3
+
+        def foo():
+            with nested_delayed(a, b, c) as (x, y, z):
+                pass
+        self.assertRaises(ZeroDivisionError, foo)
+        self.assertFalse(c_ran)
+
+class ContextManagerExitCase(unittest.TestCase):
+    """
+    Test that ContextManagerExit can be explicitl and
+    implicitly used to skip the exucution of context-managed
+    code.
+    """
+    def test_implicit(self):
+        # Test implicit creation of ContextManagerExit by skipping
+        # 'yield' in a 'contextmanager' decorated function
+        @contextmanager
+        def a():
+            return # automatically skip body
+            yield
+        def foo():
+            with a():
+                return 1
+            return 0
+        self.assertEqual(foo(), 0)
+
+    def test_explicit(self):
+        # Test the explicit raising of ContextManagerExit from __enter__()
+        class a:
+            def __enter__(self):
+                raise ContextManagerExit
+            def __exit__(self, *args):
+                pass
+        def foo():
+            with a():
+                return 1
+            return 0
+        self.assertEqual(foo(), 0)
+
+class NoneManagerCase(unittest.TestCase):
+    def test_none(self):
+        # Test that None is a valid no-op context manager
+        # that returns None as its __enter__() value.
+        with None as a:
+            self.assertEqual(a, None)
+
 class ClosingTestCase(unittest.TestCase):
 
     # XXX This needs more work
diff --git a/Objects/exceptions.c b/Objects/exceptions.c
--- a/Objects/exceptions.c
+++ b/Objects/exceptions.c
@@ -544,6 +544,12 @@
 SimpleExtendsException(PyExc_BaseException, GeneratorExit,
                        "Request that a generator exit.");
 
+/*
+ *    ContextManagerExit extends BaseException
+ */
+SimpleExtendsException(PyExc_BaseException, ContextManagerExit,
+                       "Request that a a context manager region not be enetered.");
+
 
 /*
  *    SystemExit extends BaseException
@@ -2392,6 +2398,7 @@
     PRE_INIT(TypeError)
     PRE_INIT(StopIteration)
     PRE_INIT(GeneratorExit)
+    PRE_INIT(ContextManagerExit)
     PRE_INIT(SystemExit)
     PRE_INIT(KeyboardInterrupt)
     PRE_INIT(ImportError)
@@ -2462,6 +2469,7 @@
     POST_INIT(TypeError)
     POST_INIT(StopIteration)
     POST_INIT(GeneratorExit)
+    POST_INIT(ContextManagerExit)
     POST_INIT(SystemExit)
     POST_INIT(KeyboardInterrupt)
     POST_INIT(ImportError)
diff --git a/Python/ceval.c b/Python/ceval.c
--- a/Python/ceval.c
+++ b/Python/ceval.c
@@ -2692,19 +2692,42 @@
             _Py_IDENTIFIER(__exit__);
             _Py_IDENTIFIER(__enter__);
             PyObject *mgr = TOP();
-            PyObject *exit = special_lookup(mgr, &PyId___exit__), *enter;
             PyObject *res;
-            if (exit == NULL)
-                goto error;
-            SET_TOP(exit);
-            enter = special_lookup(mgr, &PyId___enter__);
-            Py_DECREF(mgr);
-            if (enter == NULL)
-                goto error;
-            res = PyObject_CallFunctionObjArgs(enter, NULL);
-            Py_DECREF(enter);
-            if (res == NULL)
-                goto error;
+            if (mgr != Py_None) {
+                PyObject *exit = special_lookup(mgr, &PyId___exit__), *enter;
+                if (exit == NULL)
+                    goto error;
+                SET_TOP(exit);
+                enter = special_lookup(mgr, &PyId___enter__);
+                Py_DECREF(mgr);
+                if (enter == NULL)
+                    goto error;
+                res = PyObject_CallFunctionObjArgs(enter, NULL);
+                Py_DECREF(enter);
+                if (res == NULL) {
+                    if (PyErr_ExceptionMatches(PyExc_ContextManagerExit)) {
+                        /* abort the context manager and its block */
+                        PyErr_Clear();
+                        /* Replace the exit handler with null, push a None and jump to the cleanup */
+                        Py_DECREF(TOP());
+                        SET_TOP(NULL);
+                        Py_INCREF(Py_None);
+                        PUSH(Py_None);
+                        JUMPBY(oparg);
+                        DISPATCH();
+                    }
+                    goto error;
+                }
+            }
+            else
+            {
+                /* special case: A None context manager is ignored.  None is left on the
+                 * stack to be assigned to the "as" context or dropped, and
+                 * the "exit" method is left as NULL.
+                 */
+                 SET_TOP(NULL);
+                 res = Py_None;
+            }
             /* Setup the finally block before pushing the result
                of __enter__ on the stack. */
             PyFrame_BlockSetup(f, SETUP_FINALLY, INSTR_OFFSET() + oparg,
@@ -2786,25 +2809,29 @@
                 assert(block->b_type == EXCEPT_HANDLER);
                 block->b_level--;
             }
-            /* XXX Not the fastest way to call it... */
-            res = PyObject_CallFunctionObjArgs(exit_func, exc, val, tb, NULL);
-            Py_DECREF(exit_func);
-            if (res == NULL)
-                goto error;
-
-            if (exc != Py_None)
-                err = PyObject_IsTrue(res);
-            else
+            /* call the exit handler if it is non-null */
+            if (exit_func) {
+                /* XXX Not the fastest way to call it... */
+                res = PyObject_CallFunctionObjArgs(exit_func, exc, val, tb, NULL);
+                Py_DECREF(exit_func);
+                if (res == NULL)
+                    goto error;
+
+                if (exc != Py_None)
+                    err = PyObject_IsTrue(res);
+                else
+                    err = 0;
+                Py_DECREF(res);
+
+                if (err < 0)
+                    goto error;
+                else if (err > 0) {
+                    err = 0;
+                    /* There was an exception and a True return */
+                    PUSH(PyLong_FromLong((long) WHY_SILENCED));
+                }
+            } else
                 err = 0;
-            Py_DECREF(res);
-
-            if (err < 0)
-                goto error;
-            else if (err > 0) {
-                err = 0;
-                /* There was an exception and a True return */
-                PUSH(PyLong_FromLong((long) WHY_SILENCED));
-            }
             PREDICT(END_FINALLY);
             DISPATCH();
         }