# HG changeset patch # User Kristján Valur Jónsson # Date 1375957248 0 # Node ID d486faf533363236f66add94db8ff818e666ff90 # 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", "proper", "opened"] 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: @@ -115,6 +117,76 @@ return _GeneratorContextManager(func, *args, **kwds) return helper +@contextmanager +def proper(callable, *args, **kwargs): + """ + Turn a "hybrid" context manager into a proper one. A hybrid context manager is one + which allocates resources during __init__ instead of __exit__ so that its use + as a context manager is optional. An example of this is a File object. + "proper" turns this into a strict context manager by delaying the __init__ call until + __enter__ is called. This is useful if one wants to create a set of such context managers + upfront and then enter them in turn, e.g. when using "nested". + """ + with callable(*args, **kwargs) as var: + yield var + +def opened(*args, **kwargs): + return proper(open, *args, **kwargs) + +@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() + + Caution: Use this only with proper context managers that allocate resources during __enter__. + Hybrid context managares, that allocate resources during __init__ (such as files created with open) + can be durned into proper ones using the "proper" context manager: + + with nested(proper(open, filename2), proper(open, filename2)) as f1, f2: + do_something(f1, f2) + + # or use the shorthand contextlib.opened: + with nested(opened(filename2), opened(filename2)) as f1, f2: + do_something(f1, f2) + + """ + exits = [] + vars = [] + exc = (None, None, None) + try: + for mgr in managers: + 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,217 @@ 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 ProperCase(unittest.TestCase): + class Hybrid: + """ A hybrid context manager, one that allocates resource on __init__ and + releases them on __exit__ + """ + init = [0, 0] + def __init__(self, fail): + if fail: + raise ZeroDivisionError + self.init[0] += 1 + self.init[1] += 1 + def __enter__(self): + pass + def __exit__(self, *args): + self.init[0] -= 1 + + def setUp(self): + self.Hybrid.init = [0,0] + + def test_proper(self): + def foo(): + with nested(proper(self.Hybrid, False), proper(self.Hybrid, False), + proper(self.Hybrid, True), proper(self.Hybrid, True)): + pass + self.assertRaises(ZeroDivisionError, foo) + self.assertEqual(self.Hybrid.init, [0, 2]) + + def test_noproper(self): + def foo(): + with nested(self.Hybrid(False), self.Hybrid(False), + self.Hybrid(True), self.Hybrid(True)): + pass + self.assertRaises(ZeroDivisionError, foo) + self.assertEqual(self.Hybrid.init, [2, 2]) + +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(); }