Hopiu/contextlib2
1
0
Fork 0
mirror of https://github.com/jazzband/contextlib2.git synced 2026-03-16 21:50:24 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions

Issue #8: Switch to ExitStack instead of ContextStack

Browse source
This commit is contained in:
Nick Coghlan 2012-05-01 21:50:17 +10:00
parent 141ca2353f
commit 21a5d8b0a4
4 changed files with 227 additions and 66 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

2
NEWS.rst
View file

@ -4,6 +4,8 @@ Release History
0.4.0 (2012-05-??)
~~~~~~~~~~~~~~~~~~
* Issue #8: Replace ContextStack with ExitStack (old ContextStack API
retained for backwards compatibility)
* Fall back to unittest2 if unittest is missing required functionality

64
contextlib2.py
View file

@ -4,7 +4,8 @@ import sys
from collections import deque
from functools import wraps
__all__ = ["contextmanager", "closing", "ContextDecorator", "ContextStack"]
__all__ = ["contextmanager", "closing", "ContextDecorator",
"ContextStack", "ExitStack"]
class ContextDecorator(object):
@ -142,12 +143,12 @@ class closing(object):
# Inspired by discussions on http://bugs.python.org/issue13585
class ContextStack(object):
class ExitStack(object):
"""Context manager for dynamic management of a stack of exit callbacks
For example:
with ContextStack() as stack:
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
@ -155,23 +156,23 @@ class ContextStack(object):
"""
def __init__(self):
self._callbacks = deque()
self._exit_callbacks = deque()
def preserve(self):
def pop_all(self):
"""Preserve the context stack by transferring it to a new instance"""
new_stack = type(self)()
new_stack._callbacks = self._callbacks
self._callbacks = deque()
new_stack._exit_callbacks = self._exit_callbacks
self._exit_callbacks = deque()
return new_stack
def _register_cm_exit(self, cm, cm_exit):
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.register_exit(_exit_wrapper)
self.push(_exit_wrapper)
def register_exit(self, callback):
def push(self, exit):
"""Registers a callback with the standard __exit__ method signature
Can suppress exceptions the same way __exit__ methods can.
@ -179,16 +180,19 @@ class ContextStack(object):
Also accepts any object with an __exit__ method (registering the
method instead of the object itself)
"""
_cb_type = type(callback)
# 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 = _cb_type.__exit__
exit_method = _cb_type.__exit__
except AttributeError:
self._callbacks.append(callback)
# Not a context manager, so assume its a callable
self._exit_callbacks.append(exit)
else:
self._register_cm_exit(callback, exit)
return callback # Allow use as a decorator
self._push_cm_exit(exit, exit_method)
return exit # Allow use as a decorator
def register(self, callback, *args, **kwds):
def callback(self, callback, *args, **kwds):
"""Registers an arbitrary callback and arguments.
Cannot suppress exceptions.
@ -198,19 +202,20 @@ class ContextStack(object):
# We changed the signature, so using @wraps is not appropriate, but
# setting __wrapped__ may still help with introspection
_exit_wrapper.__wrapped__ = callback
self.register_exit(_exit_wrapper)
self.push(_exit_wrapper)
return callback # Allow use as a decorator
def enter_context(self, cm):
"""Enters the supplied context manager
If successful, also registers its __exit__ method as a callback and
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._register_cm_exit(cm, _exit)
self._push_cm_exit(cm, _exit)
return result
def close(self):
@ -221,7 +226,7 @@ class ContextStack(object):
return self
def __exit__(self, *exc_details):
if not self._callbacks:
if not self._exit_callbacks:
return
# This looks complicated, but it is really just
# setting up a chain of try-expect statements to ensure
@ -230,8 +235,8 @@ class ContextStack(object):
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._callbacks.popleft()
if not self._callbacks:
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
@ -249,4 +254,17 @@ class ContextStack(object):
suppress_exc = cb(*exc_details) or suppress_exc
return suppress_exc
# Kick off the recursive chain
return _invoke_next_callback(exc_details)
return _invoke_next_callback(exc_details)
# Preserve backwards compatibility
class ContextStack(ExitStack):
"""Backwards compatibility alias for ExitStack"""
def register_exit(self, callback):
return self.push(callback)
def register(self, callback, *args, **kwds):
return self.callback(callback, *args, **kwds)
def preserve(self):
return self.pop_all()

74
docs/index.rst
View file

@ -26,7 +26,7 @@ This module is primarily a backport of the Python 3.2 version of
for new features not yet part of the standard library. Those new features
are currently:
* :class:`ContextStack`
* :class:`ExitStack`
* :meth:`ContextDecorator.refresh_cm`
@ -192,7 +192,7 @@ API Reference
Made the standard library's private :meth:`refresh_cm` API public
.. class:: ContextStack()
.. class:: ExitStack()
A context manager that is designed to make it easy to programmatically
combine other context managers and cleanup functions, especially those
@ -201,7 +201,7 @@ API Reference
For example, a set of files may easily be handled in a single with
statement as follows::
with ContextStack() as stack:
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
@ -228,7 +228,7 @@ API Reference
These context managers may suppress exceptions just as they normally
would if used directly as part of a ``with`` statement.
.. method:: register_exit(callback)
.. method:: push(exit)
Directly accepts a callback with the same signature as a
context manager's :meth:`__exit__` method and adds it to the callback
@ -242,7 +242,7 @@ API Reference
cover part of an :meth:`__enter__` implementation with a context
manager's own :meth:`__exit__` method.
.. method:: register(callback, *args, **kwds)
.. method:: callback(callback, *args, **kwds)
Accepts an arbitrary callback function and arguments and adds it to
the callback stack.
@ -250,7 +250,7 @@ API Reference
Unlike the other methods, callbacks added this way cannot suppress
exceptions (as they are never passed the exception details).
.. method:: preserve()
.. method:: pop_all()
Transfers the callback stack to a fresh instance and returns it. No
callbacks are invoked by this operation - instead, they will now be
@ -259,16 +259,14 @@ API Reference
For example, a group of files can be opened as an "all or nothing"
operation as follows::
with ContextStack() as stack:
with ExitStack() as stack:
files = [stack.enter_context(open(fname)) for fname in filenames]
close_files = stack.preserve().close
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
.. versionadded:: 0.3
.. method:: close()
Immediately unwinds the callback stack, invoking callbacks in the
@ -276,6 +274,18 @@ API Reference
callbacks registered, the arguments passed in will indicate that no
exception occurred.
.. versionadded:: 0.4
New API for :mod:`contextlib2`, not available in standard library
.. class:: ContextStack()
An earlier incarnation of the :class:`ExitStack` interface. This class
is deprecated and should no longer be used.
.. versionchanged:: 0.4
Deprecated in favour of :class:`ExitStack`
.. versionadded:: 0.2
New API for :mod:`contextlib2`, not available in standard library
@ -339,7 +349,7 @@ This example should also work with :mod:`contextlib` in Python 3.2.1 or later.
Cleaning up in an ``__enter__`` implementation
----------------------------------------------
As noted in the documentation of :meth:`ContextStack.register_exit`, this
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.
@ -347,7 +357,7 @@ 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 contextlib2 import ContextStack
from contextlib2 import ExitStack
class ResourceManager(object):
@ -359,15 +369,15 @@ and maps them to the context management protocol::
def __enter__(self):
resource = self.acquire_resource()
if self.check_resource_ok is not None:
with ContextStack() as stack:
stack.register_exit(self)
with ExitStack() as stack:
stack.push(self)
if not self.check_resource_ok(resource):
msg = "Failed validation for {!r}"
raise RuntimeError(msg.format(resource))
# 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.preserve()
stack.pop_all()
return resource
def __exit__(self, *exc_details):
@ -396,17 +406,17 @@ 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:`ContextStack` makes it possible to instead register a callback for
: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 contextlib2 import ContextStack
from contextlib2 import ExitStack
with ContextStack() as stack:
stack.register(cleanup_resources)
with ExitStack() as stack:
stack.callback(cleanup_resources)
result = perform_operation()
if result:
stack.preserve()
stack.pop_all()
This allows the intended cleanup up behaviour to be made explicit up front,
rather than requiring a separate flag variable.
@ -414,15 +424,15 @@ rather than requiring a separate flag variable.
If you find yourself using this pattern a lot, it can be simplified even
further by means of a small helper class::
from contextlib2 import ContextStack
from contextlib2 import ExitStack
class Callback(ContextStack):
class Callback(ExitStack):
def __init__(self, callback, *args, **kwds):
super(Callback, self).__init__()
self.register(callback, *args, **kwds)
self.callback(callback, *args, **kwds)
def cancel(self):
self.preserve()
self.pop_all()
with Callback(cleanup_resources) as cb:
result = perform_operation()
@ -431,18 +441,22 @@ further by means of a small helper class::
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:`ContextStack.register_exit` to declare the resource cleanup in
:meth:`ExitStack.callback` to declare the resource cleanup in
advance::
from contextlib2 import ContextStack
from contextlib2 import ExitStack
with ContextStack() as stack:
@stack.register_exit
def cleanup_resources(*exc_details):
with ExitStack() as stack:
@stack.callback
def cleanup_resources():
...
result = perform_operation()
if result:
stack.preserve()
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
Obtaining the Module

153
test_contextlib2.py
View file

@ -303,6 +303,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)
class TestContextStack(unittest.TestCase):
def test_no_resources(self):
@ -325,14 +448,15 @@ class TestContextStack(unittest.TestCase):
with ContextStack() as stack:
for args, kwds in reversed(expected):
if args and kwds:
self.assertIsNone(stack.register(_exit, *args, **kwds))
f = stack.register(_exit, *args, **kwds)
elif args:
self.assertIsNone(stack.register(_exit, *args))
f = stack.register(_exit, *args)
elif kwds:
self.assertIsNone(stack.register(_exit, **kwds))
f = stack.register(_exit, **kwds)
else:
self.assertIsNone(stack.register(_exit))
for wrapper in stack._callbacks:
f = stack.register(_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__)
@ -357,19 +481,19 @@ class TestContextStack(unittest.TestCase):
self.check_exc(*exc_details)
with ContextStack() as stack:
stack.register_exit(_expect_ok)
self.assertIs(stack._callbacks[-1], _expect_ok)
self.assertIs(stack._exit_callbacks[-1], _expect_ok)
cm = ExitCM(_expect_ok)
stack.register_exit(cm)
self.assertIs(stack._callbacks[-1].__self__, cm)
self.assertIs(stack._exit_callbacks[-1].__self__, cm)
stack.register_exit(_suppress_exc)
self.assertIs(stack._callbacks[-1], _suppress_exc)
self.assertIs(stack._exit_callbacks[-1], _suppress_exc)
cm = ExitCM(_expect_exc)
stack.register_exit(cm)
self.assertIs(stack._callbacks[-1].__self__, cm)
self.assertIs(stack._exit_callbacks[-1].__self__, cm)
stack.register_exit(_expect_exc)
self.assertIs(stack._callbacks[-1], _expect_exc)
self.assertIs(stack._exit_callbacks[-1], _expect_exc)
stack.register_exit(_expect_exc)
self.assertIs(stack._callbacks[-1], _expect_exc)
self.assertIs(stack._exit_callbacks[-1], _expect_exc)
1/0
def test_enter_context(self):
@ -385,8 +509,9 @@ class TestContextStack(unittest.TestCase):
@stack.register # Registered first => cleaned up last
def _exit():
result.append(4)
self.assertIsNotNone(_exit)
stack.enter_context(cm)
self.assertIs(stack._callbacks[-1].__self__, cm)
self.assertIs(stack._exit_callbacks[-1].__self__, cm)
result.append(2)
self.assertEqual(result, [1, 2, 3, 4])
@ -396,6 +521,7 @@ class TestContextStack(unittest.TestCase):
@stack.register
def _exit():
result.append(1)
self.assertIsNotNone(_exit)
stack.close()
result.append(2)
self.assertEqual(result, [1, 2])
@ -406,6 +532,7 @@ class TestContextStack(unittest.TestCase):
@stack.register
def _exit():
result.append(3)
self.assertIsNotNone(_exit)
new_stack = stack.preserve()
result.append(1)
result.append(2)
@ -419,7 +546,7 @@ class TestContextStack(unittest.TestCase):
stack = ContextStack()
self.assertRaises(AttributeError, stack.enter_context, cm)
stack.register_exit(cm)
self.assertIs(stack._callbacks[-1], cm)
self.assertIs(stack._exit_callbacks[-1], cm)
if __name__ == "__main__":
import unittest