From 3267a30de12724971bd7e7b9262e19c1d9b2becd Mon Sep 17 00:00:00 2001 From: Nick Coghlan Date: Mon, 21 May 2012 22:54:43 +1000 Subject: Close #13585: add contextlib.ExitStack to replace the ill-fated contextlib.nested API --- Lib/contextlib.py | 126 +++++++++++++++++++++++++++++++++++++++++++++++++++--- 1 file changed, 121 insertions(+), 5 deletions(-) (limited to 'Lib/contextlib.py') diff --git a/Lib/contextlib.py b/Lib/contextlib.py index 2f8f00d2964..ead11554cb1 100644 --- a/Lib/contextlib.py +++ b/Lib/contextlib.py @@ -1,9 +1,10 @@ """Utilities for with-statement contexts. See PEP 343.""" import sys +from collections import deque from functools import wraps -__all__ = ["contextmanager", "closing", "ContextDecorator"] +__all__ = ["contextmanager", "closing", "ContextDecorator", "ExitStack"] class ContextDecorator(object): @@ -12,12 +13,12 @@ class ContextDecorator(object): def _recreate_cm(self): """Return a recreated instance of self. - Allows otherwise one-shot context managers like + Allows an otherwise one-shot context manager like _GeneratorContextManager to support use as - decorators via implicit recreation. + a decorator via implicit recreation. - Note: this is a private interface just for _GCM in 3.2 but will be - renamed and documented for third party use in 3.3 + This is a private interface just for _GeneratorContextManager. + See issue #11647 for details. """ return self @@ -138,3 +139,118 @@ def __enter__(self): return self.thing def __exit__(self, *exc_info): self.thing.close() + + +# Inspired by discussions on http://bugs.python.org/issue13585 +class ExitStack(object): + """Context manager for dynamic management of a stack of exit callbacks + + For example: + + with ExitStack() as stack: + files = [stack.enter_context(open(fname)) for fname in filenames] + # All opened files will automatically be closed at the end of + # the with statement, even if attempts to open files later + # in the list throw an exception + + """ + def __init__(self): + self._exit_callbacks = deque() + + def pop_all(self): + """Preserve the context stack by transferring it to a new instance""" + new_stack = type(self)() + new_stack._exit_callbacks = self._exit_callbacks + self._exit_callbacks = deque() + return new_stack + + def _push_cm_exit(self, cm, cm_exit): + """Helper to correctly register callbacks to __exit__ methods""" + def _exit_wrapper(*exc_details): + return cm_exit(cm, *exc_details) + _exit_wrapper.__self__ = cm + self.push(_exit_wrapper) + + def push(self, exit): + """Registers a callback with the standard __exit__ method signature + + Can suppress exceptions the same way __exit__ methods can. + + Also accepts any object with an __exit__ method (registering a call + to the method instead of the object itself) + """ + # We use an unbound method rather than a bound method to follow + # the standard lookup behaviour for special methods + _cb_type = type(exit) + try: + exit_method = _cb_type.__exit__ + except AttributeError: + # Not a context manager, so assume its a callable + self._exit_callbacks.append(exit) + else: + self._push_cm_exit(exit, exit_method) + return exit # Allow use as a decorator + + def callback(self, callback, *args, **kwds): + """Registers an arbitrary callback and arguments. + + Cannot suppress exceptions. + """ + def _exit_wrapper(exc_type, exc, tb): + callback(*args, **kwds) + # We changed the signature, so using @wraps is not appropriate, but + # setting __wrapped__ may still help with introspection + _exit_wrapper.__wrapped__ = callback + self.push(_exit_wrapper) + return callback # Allow use as a decorator + + def enter_context(self, cm): + """Enters the supplied context manager + + If successful, also pushes its __exit__ method as a callback and + returns the result of the __enter__ method. + """ + # We look up the special methods on the type to match the with statement + _cm_type = type(cm) + _exit = _cm_type.__exit__ + result = _cm_type.__enter__(cm) + self._push_cm_exit(cm, _exit) + return result + + def close(self): + """Immediately unwind the context stack""" + self.__exit__(None, None, None) + + def __enter__(self): + return self + + def __exit__(self, *exc_details): + if not self._exit_callbacks: + return + # This looks complicated, but it is really just + # setting up a chain of try-expect statements to ensure + # that outer callbacks still get invoked even if an + # inner one throws an exception + def _invoke_next_callback(exc_details): + # Callbacks are removed from the list in FIFO order + # but the recursion means they're invoked in LIFO order + cb = self._exit_callbacks.popleft() + if not self._exit_callbacks: + # Innermost callback is invoked directly + return cb(*exc_details) + # More callbacks left, so descend another level in the stack + try: + suppress_exc = _invoke_next_callback(exc_details) + except: + suppress_exc = cb(*sys.exc_info()) + # Check if this cb suppressed the inner exception + if not suppress_exc: + raise + else: + # Check if inner cb suppressed the original exception + if suppress_exc: + exc_details = (None, None, None) + suppress_exc = cb(*exc_details) or suppress_exc + return suppress_exc + # Kick off the recursive chain + return _invoke_next_callback(exc_details) -- cgit v1.2.3-65-gdbad