Close #13585: add contextlib.ExitStack to replace the ill-fated contextlib.nested API
This commit is contained in:
Nick Coghlan
parent
6e49ac2406
commit
3267a30de1
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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
|
||||
-----
|
||||
|
||||
|
|
|
|||
|
|
@ -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)
|
||||
|
|
|
|||
|
|
@ -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__)
|
||||
|
|
|
|||
Loading…
Reference in New Issue