contextlib: port ExitStack from CPython 3.4

contextlib/contextlib.py

@@ -4,10 +4,11 @@
 
 Not implemented:
  - redirect_stdout;
- - ExitStack.
 
 """
 
+import sys
+from collections import deque
 from ucontextlib import *
 
 
@@ -64,3 +65,102 @@ def __exit__(self, exctype, excinst, exctb):
         #
         # See http://bugs.python.org/issue12029 for more details
         return exctype is not None and issubclass(exctype, self._exceptions)
+
+# 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 raise 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)
+        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)
+        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):
+        received_exc = exc_details[0] is not None
+        # Callbacks are invoked in LIFO order to match the behaviour of
+        # nested context managers
+        suppressed_exc = False
+        pending_raise = False
+        while self._exit_callbacks:
+            cb = self._exit_callbacks.pop()
+            try:
+                if cb(*exc_details):
+                    suppressed_exc = True
+                    pending_raise = False
+                    exc_details = (None, None, None)
+            except:
+                exc_details = sys.exc_info()
+                pending_raise = True
+        if pending_raise:
+            raise exc_details[1]
+        return received_exc and suppressed_exc

contextlib/metadata.txt

@@ -1,5 +1,5 @@
 srctype = cpython
 type = module
-version = 3.4.2-2
+version = 3.4.2-3
 long_desc = Port of contextlib for micropython
-depends = ucontextlib
+depends = ucontextlib collections

contextlib/setup.py

@@ -6,7 +6,7 @@
 
 
 setup(name='micropython-contextlib',
-      version='3.4.2-2',
+      version='3.4.2-3',
       description='CPython contextlib module ported to MicroPython',
       long_description='This is a module ported from CPython standard library to be compatible with\nMicroPython interpreter. Usually, this means applying small patches for\nfeatures not supported (yet, or at all) in MicroPython. Sometimes, heavier\nchanges are required. Note that CPython modules are written with availability\nof vast resources in mind, and may not work for MicroPython ports with\nlimited heap. If you are affected by such a case, please help reimplement\nthe module from scratch.',
       url='https://github.com/micropython/micropython/issues/405',
@@ -16,4 +16,4 @@
       maintainer_email='micro-python@googlegroups.com',
       license='Python',
       py_modules=['contextlib'],
-      install_requires=['micropython-ucontextlib'])
+      install_requires=['micropython-ucontextlib collections'])