From 3267a30de12724971bd7e7b9262e19c1d9b2becd Mon Sep 17 00:00:00 2001 From: Nick Coghlan Date: Mon, 21 May 2012 22:54:43 +1000 Subject: [PATCH] Close #13585: add contextlib.ExitStack to replace the ill-fated contextlib.nested API --- Doc/library/contextlib.rst | 279 +++++++++++++++++++++++++++++++++++- Doc/whatsnew/3.3.rst | 15 ++ Lib/contextlib.py | 126 +++++++++++++++- Lib/test/test_contextlib.py | 123 ++++++++++++++++ Misc/NEWS | 2 + 5 files changed, 539 insertions(+), 6 deletions(-) diff --git a/Doc/library/contextlib.rst b/Doc/library/contextlib.rst index e8dc17fed32..8cc71b4f399 100644 --- a/Doc/library/contextlib.rst +++ b/Doc/library/contextlib.rst @@ -12,8 +12,11 @@ This module provides utilities for common tasks involving the :keyword:`with` statement. For more information see also :ref:`typecontextmanager` and :ref:`context-managers`. -Functions provided: +Utilities +--------- + +Functions and classes provided: .. decorator:: contextmanager @@ -168,6 +171,280 @@ Functions provided: .. versionadded:: 3.2 +.. class:: ExitStack() + + A context manager that is designed to make it easy to programmatically + combine other context managers and cleanup functions, especially those + that are optional or otherwise driven by input data. + + For example, a set of files may easily be handled in a single with + statement as follows:: + + 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 + + Each instance maintains a stack of registered callbacks that are called in + reverse order when the instance is closed (either explicitly or implicitly + at the end of a ``with`` statement). Note that callbacks are *not* invoked + implicitly when the context stack instance is garbage collected. + + This stack model is used so that context managers that acquire their + resources in their ``__init__`` method (such as file objects) can be + handled correctly. + + Since registered callbacks are invoked in the reverse order of + registration, this ends up behaving as if multiple nested ``with`` + statements had been used with the registered set of callbacks. This even + extends to exception handling - if an inner callback suppresses or replaces + an exception, then outer callbacks will be passed arguments based on that + updated state. + + This is a relatively low level API that takes care of the details of + correctly unwinding the stack of exit callbacks. It provides a suitable + foundation for higher level context managers that manipulate the exit + stack in application specific ways. + + .. method:: enter_context(cm) + + Enters a new context manager and adds its :meth:`__exit__` method to + the callback stack. The return value is the result of the context + manager's own :meth:`__enter__` method. + + These context managers may suppress exceptions just as they normally + would if used directly as part of a ``with`` statement. + + .. method:: push(exit) + + Adds a context manager's :meth:`__exit__` method to the callback stack. + + As ``__enter__`` is *not* invoked, this method can be used to cover + part of an :meth:`__enter__` implementation with a context manager's own + :meth:`__exit__` method. + + If passed an object that is not a context manager, this method assumes + it is a callback with the same signature as a context manager's + :meth:`__exit__` method and adds it directly to the callback stack. + + By returning true values, these callbacks can suppress exceptions the + same way context manager :meth:`__exit__` methods can. + + The passed in object is returned from the function, allowing this + method to be used is a function decorator. + + .. method:: callback(callback, *args, **kwds) + + Accepts an arbitrary callback function and arguments and adds it to + the callback stack. + + Unlike the other methods, callbacks added this way cannot suppress + exceptions (as they are never passed the exception details). + + The passed in callback is returned from the function, allowing this + method to be used is a function decorator. + + .. method:: pop_all() + + Transfers the callback stack to a fresh :class:`ExitStack` instance + and returns it. No callbacks are invoked by this operation - instead, + they will now be invoked when the new stack is closed (either + explicitly or implicitly). + + For example, a group of files can be opened as an "all or nothing" + operation as follows:: + + with ExitStack() as stack: + files = [stack.enter_context(open(fname)) for fname in filenames] + close_files = stack.pop_all().close + # If opening any file fails, all previously opened files will be + # closed automatically. If all files are opened successfully, + # they will remain open even after the with statement ends. + # close_files() can then be invoked explicitly to close them all + + .. method:: close() + + Immediately unwinds the callback stack, invoking callbacks in the + reverse order of registration. For any context managers and exit + callbacks registered, the arguments passed in will indicate that no + exception occurred. + + .. versionadded:: 3.3 + + +Examples and Recipes +-------------------- + +This section describes some examples and recipes for making effective use of +the tools provided by :mod:`contextlib`. + + +Cleaning up in an ``__enter__`` implementation +^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ + +As noted in the documentation of :meth:`ExitStack.push`, this +method can be useful in cleaning up an already allocated resource if later +steps in the :meth:`__enter__` implementation fail. + +Here's an example of doing this for a context manager that accepts resource +acquisition and release functions, along with an optional validation function, +and maps them to the context management protocol:: + + from contextlib import contextmanager, ExitStack + + class ResourceManager(object): + + def __init__(self, acquire_resource, release_resource, check_resource_ok=None): + self.acquire_resource = acquire_resource + self.release_resource = release_resource + if check_resource_ok is None: + def check_resource_ok(resource): + return True + self.check_resource_ok = check_resource_ok + + @contextmanager + def _cleanup_on_error(self): + with ExitStack() as stack: + stack.push(self) + yield + # The validation check passed and didn't raise an exception + # Accordingly, we want to keep the resource, and pass it + # back to our caller + stack.pop_all() + + def __enter__(self): + resource = self.acquire_resource() + with self._cleanup_on_error(): + if not self.check_resource_ok(resource): + msg = "Failed validation for {!r}" + raise RuntimeError(msg.format(resource)) + return resource + + def __exit__(self, *exc_details): + # We don't need to duplicate any of our resource release logic + self.release_resource() + + +Replacing any use of ``try-finally`` and flag variables +^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ + +A pattern you will sometimes see is a ``try-finally`` statement with a flag +variable to indicate whether or not the body of the ``finally`` clause should +be executed. In its simplest form (that can't already be handled just by +using an ``except`` clause instead), it looks something like this:: + + cleanup_needed = True + try: + result = perform_operation() + if result: + cleanup_needed = False + finally: + if cleanup_needed: + cleanup_resources() + +As with any ``try`` statement based code, this can cause problems for +development and review, because the setup code and the cleanup code can end +up being separated by arbitrarily long sections of code. + +:class:`ExitStack` makes it possible to instead register a callback for +execution at the end of a ``with`` statement, and then later decide to skip +executing that callback:: + + from contextlib import ExitStack + + with ExitStack() as stack: + stack.callback(cleanup_resources) + result = perform_operation() + if result: + stack.pop_all() + +This allows the intended cleanup up behaviour to be made explicit up front, +rather than requiring a separate flag variable. + +If a particular application uses this pattern a lot, it can be simplified +even further by means of a small helper class:: + + from contextlib import ExitStack + + class Callback(ExitStack): + def __init__(self, callback, *args, **kwds): + super(Callback, self).__init__() + self.callback(callback, *args, **kwds) + + def cancel(self): + self.pop_all() + + with Callback(cleanup_resources) as cb: + result = perform_operation() + if result: + cb.cancel() + +If the resource cleanup isn't already neatly bundled into a standalone +function, then it is still possible to use the decorator form of +:meth:`ExitStack.callback` to declare the resource cleanup in +advance:: + + from contextlib import ExitStack + + with ExitStack() as stack: + @stack.callback + def cleanup_resources(): + ... + result = perform_operation() + if result: + stack.pop_all() + +Due to the way the decorator protocol works, a callback function +declared this way cannot take any parameters. Instead, any resources to +be released must be accessed as closure variables + + +Using a context manager as a function decorator +^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ + +:class:`ContextDecorator` makes it possible to use a context manager in +both an ordinary ``with`` statement and also as a function decorator. + +For example, it is sometimes useful to wrap functions or groups of statements +with a logger that can track the time of entry and time of exit. Rather than +writing both a function decorator and a context manager for the task, +inheriting from :class:`ContextDecorator` provides both capabilities in a +single definition:: + + from contextlib import ContextDecorator + import logging + + logging.basicConfig(level=logging.INFO) + + class track_entry_and_exit(ContextDecorator): + def __init__(self, name): + self.name = name + + def __enter__(self): + logging.info('Entering: {}'.format(name)) + + def __exit__(self, exc_type, exc, exc_tb): + logging.info('Exiting: {}'.format(name)) + +Instances of this class can be used as both a context manager:: + + with track_entry_and_exit('widget loader'): + print('Some time consuming activity goes here') + load_widget() + +And also as a function decorator:: + + @track_entry_and_exit('widget loader') + def activity(): + print('Some time consuming activity goes here') + load_widget() + +Note that there is one additional limitation when using context managers +as function decorators: there's no way to access the return value of +:meth:`__enter__`. If that value is needed, then it is still necessary to use +an explicit ``with`` statement. + .. seealso:: :pep:`0343` - The "with" statement diff --git a/Doc/whatsnew/3.3.rst b/Doc/whatsnew/3.3.rst index 9afd8ae3efb..e8d9d98c7a3 100644 --- a/Doc/whatsnew/3.3.rst +++ b/Doc/whatsnew/3.3.rst @@ -696,6 +696,21 @@ collections classes. Aliases for ABCs are still present in the .. XXX addition of __slots__ to ABCs not recorded here: internal detail +contextlib +---------- + +:class:`~collections.ExitStack` now provides a solid foundation for +programmatic manipulation of context managers and similar cleanup +functionality. Unlike the previous ``contextlib.nested`` API (which was +deprecated and removed), the new API is designed to work correctly +regardless of whether context managers acquire their resources in +their ``__init`` method (for example, file objects) or in their +``__enter__`` method (for example, synchronisation objects from the +:mod:`threading` module). + +(:issue:`13585`) + + crypt ----- 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 @@ 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): + 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) diff --git a/Lib/test/test_contextlib.py b/Lib/test/test_contextlib.py index 6e38305123d..8bed88e8ca4 100644 --- a/Lib/test/test_contextlib.py +++ b/Lib/test/test_contextlib.py @@ -370,6 +370,129 @@ class TestContextDecorator(unittest.TestCase): self.assertEqual(state, [1, 'something else', 999]) +class TestExitStack(unittest.TestCase): + + def test_no_resources(self): + with ExitStack(): + pass + + def test_callback(self): + expected = [ + ((), {}), + ((1,), {}), + ((1,2), {}), + ((), dict(example=1)), + ((1,), dict(example=1)), + ((1,2), dict(example=1)), + ] + result = [] + def _exit(*args, **kwds): + """Test metadata propagation""" + result.append((args, kwds)) + with ExitStack() as stack: + for args, kwds in reversed(expected): + if args and kwds: + f = stack.callback(_exit, *args, **kwds) + elif args: + f = stack.callback(_exit, *args) + elif kwds: + f = stack.callback(_exit, **kwds) + else: + f = stack.callback(_exit) + self.assertIs(f, _exit) + for wrapper in stack._exit_callbacks: + self.assertIs(wrapper.__wrapped__, _exit) + self.assertNotEqual(wrapper.__name__, _exit.__name__) + self.assertIsNone(wrapper.__doc__, _exit.__doc__) + self.assertEqual(result, expected) + + def test_push(self): + exc_raised = ZeroDivisionError + def _expect_exc(exc_type, exc, exc_tb): + self.assertIs(exc_type, exc_raised) + def _suppress_exc(*exc_details): + return True + def _expect_ok(exc_type, exc, exc_tb): + 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 ExitStack() as stack: + stack.push(_expect_ok) + self.assertIs(stack._exit_callbacks[-1], _expect_ok) + cm = ExitCM(_expect_ok) + stack.push(cm) + self.assertIs(stack._exit_callbacks[-1].__self__, cm) + stack.push(_suppress_exc) + self.assertIs(stack._exit_callbacks[-1], _suppress_exc) + cm = ExitCM(_expect_exc) + stack.push(cm) + self.assertIs(stack._exit_callbacks[-1].__self__, cm) + stack.push(_expect_exc) + self.assertIs(stack._exit_callbacks[-1], _expect_exc) + stack.push(_expect_exc) + self.assertIs(stack._exit_callbacks[-1], _expect_exc) + 1/0 + + def test_enter_context(self): + class TestCM(object): + def __enter__(self): + result.append(1) + def __exit__(self, *exc_details): + result.append(3) + + result = [] + cm = TestCM() + with ExitStack() as stack: + @stack.callback # Registered first => cleaned up last + def _exit(): + result.append(4) + self.assertIsNotNone(_exit) + stack.enter_context(cm) + self.assertIs(stack._exit_callbacks[-1].__self__, cm) + result.append(2) + self.assertEqual(result, [1, 2, 3, 4]) + + def test_close(self): + result = [] + with ExitStack() as stack: + @stack.callback + def _exit(): + result.append(1) + self.assertIsNotNone(_exit) + stack.close() + result.append(2) + self.assertEqual(result, [1, 2]) + + def test_pop_all(self): + result = [] + with ExitStack() as stack: + @stack.callback + def _exit(): + result.append(3) + self.assertIsNotNone(_exit) + new_stack = stack.pop_all() + result.append(1) + result.append(2) + new_stack.close() + self.assertEqual(result, [1, 2, 3]) + + def test_instance_bypass(self): + class Example(object): pass + cm = Example() + cm.__exit__ = object() + stack = ExitStack() + self.assertRaises(AttributeError, stack.enter_context, cm) + stack.push(cm) + self.assertIs(stack._exit_callbacks[-1], cm) + + # This is needed to make the test actually run under regrtest.py! def test_main(): support.run_unittest(__name__) diff --git a/Misc/NEWS b/Misc/NEWS index 699cf7ecfd3..ae4f7aab809 100644 --- a/Misc/NEWS +++ b/Misc/NEWS @@ -42,6 +42,8 @@ Core and Builtins Library ------- +- Issue #13585: Added contextlib.ExitStack + - PEP 3144, Issue #14814: Added the ipaddress module - Issue #14426: Correct the Date format in Expires attribute of Set-Cookie