summaryrefslogtreecommitdiffstats
path: root/Lib/contextlib.py
diff options
context:
space:
mode:
Diffstat (limited to 'Lib/contextlib.py')
-rw-r--r--Lib/contextlib.py126
1 files changed, 120 insertions, 6 deletions
diff --git a/Lib/contextlib.py b/Lib/contextlib.py
index 5ebbbc6..bde2feb 100644
--- a/Lib/contextlib.py
+++ b/Lib/contextlib.py
@@ -1,10 +1,10 @@
"""Utilities for with-statement contexts. See PEP 343."""
import sys
+from collections import deque
from functools import wraps
-from warnings import warn
-__all__ = ["contextmanager", "closing", "ContextDecorator"]
+__all__ = ["contextmanager", "closing", "ContextDecorator", "ExitStack"]
class ContextDecorator(object):
@@ -13,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
@@ -139,3 +139,117 @@ class closing(object):
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):
+ # We manipulate the exception state so it behaves as though
+ # we were actually nesting multiple with statements
+ frame_exc = sys.exc_info()[1]
+ def _fix_exception_context(new_exc, old_exc):
+ while 1:
+ exc_context = new_exc.__context__
+ if exc_context in (None, frame_exc):
+ break
+ new_exc = exc_context
+ new_exc.__context__ = old_exc
+
+ # Callbacks are invoked in LIFO order to match the behaviour of
+ # nested context managers
+ suppressed_exc = False
+ while self._exit_callbacks:
+ cb = self._exit_callbacks.pop()
+ try:
+ if cb(*exc_details):
+ suppressed_exc = True
+ exc_details = (None, None, None)
+ except:
+ new_exc_details = sys.exc_info()
+ # simulate the stack of exceptions by setting the context
+ _fix_exception_context(new_exc_details[1], exc_details[1])
+ if not self._exit_callbacks:
+ raise
+ exc_details = new_exc_details
+ return suppressed_exc