diff --git a/NEWS.rst b/NEWS.rst index 7ff5ef6..4834488 100644 --- a/NEWS.rst +++ b/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) ~~~~~~~~~~~~~~~~ diff --git a/contextlib2.py b/contextlib2.py index 6dbd7ce..2b95040 100644 --- a/contextlib2.py +++ b/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): diff --git a/docs/index.rst b/docs/index.rst index 91daf82..b0111c3 100644 --- a/docs/index.rst +++ b/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 diff --git a/test_contextlib2.py b/test_contextlib2.py index db1caa2..31b95b0 100755 --- a/test_contextlib2.py +++ b/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