Issue #3: Accept context managers in ContextStack.register_exit()

NEWS.rst

@@ -5,18 +5,22 @@ Release History
 0.3 (2012-01-XX)
 ~~~~~~~~~~~~~~~~
 
+* Issue #3: ContextStack.register_exit() now accepts objects with __exit__
+  attributes in addition to accepting exit callbacks directly
 * Issue #1: Add ContextStack.preserve() to move all registered callbacks to
   a new ContextStack object
 * Wrapped callbacks now use functools.wraps to aid in introspection
 * Moved version number to a VERSION.txt file (read by both docs and setup.py)
 * Added NEWS.rst (and incorporated into documentation)
 
+
 0.2 (2011-12-15)
 ~~~~~~~~~~~~~~~~
 
 * Renamed CleanupManager to ContextStack (hopefully before anyone started
   using the module for anything, since I didn't alias the old name at all)
 
+
 0.1 (2011-12-13)
 ~~~~~~~~~~~~~~~~
 

contextlib2.py

@@ -168,8 +168,15 @@ class ContextStack(object):
         """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 the
+        method instead of the object itself)
         """
-        self._callbacks.append(callback)
+        try:
+            exit = callback.__exit__
+        except AttributeError:
+            exit = callback
+        self._callbacks.append(exit)
         return callback # Allow use as a decorator
 
     def register(self, callback, *args, **kwds):

docs/index.rst

@@ -237,6 +237,11 @@ API Reference
       By returning true values, these callbacks can suppress exceptions the
       same way context manager :meth:`__exit__` methods can.
 
+      This method also accepts any object with an ``__exit__`` method, and
+      will register that method as the callback. This is mainly useful to
+      cover part of an :meth:`__enter__` implementation with a context
+      manager's own :meth:`__exit__` method.
+
    .. method:: register(callback, *args, **kwds)
 
       Accepts an arbitrary callback function and arguments and adds it to

test_contextlib2.py

@@ -344,9 +344,18 @@ class TestContextStack(unittest.TestCase):
             self.assertIsNone(exc_type)
             self.assertIsNone(exc)
             self.assertIsNone(exc_tb)
+        class ExitCM(object):
+            def __init__(self, check_exc):
+                self.check_exc = check_exc
+            def __enter__(self):
+                self.fail("Should not be called!")
+            def __exit__(self, *exc_details):
+                self.check_exc(*exc_details)
         with ContextStack() as stack:
             stack.register_exit(_expect_ok)
+            stack.register_exit(ExitCM(_expect_ok))
             stack.register_exit(_suppress_exc)
+            stack.register_exit(ExitCM(_expect_exc))
             stack.register_exit(_expect_exc)
             stack.register_exit(_expect_exc)
             1/0