Welcome, guest | Sign In | My Account | Store | Cart

Inspired by unittest.TestCase.addCleanup(), CleanupManager provides a means to programmatically add resources to be cleaned up when leaving a with statement. This makes it easy to use with optional resources, and those derived from sequences of inputs.

An more powerful version of this recipe with a few additional features is published under the name ContextStack as part of the contextlib2 module: http://contextlib2.readthedocs.org

This recipe is based on a suggestion originally posted by Nikolaus Rath at http://bugs.python.org/issue13585

Python, 68 lines
 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
import collections
import sys

class CleanupManager(object):
    """Programmatic management of resource cleanup"""
    def __init__(self):
        self._callbacks = collections.deque()

    def register_exit(self, exit):
        """Accepts callbacks with the same signature as context manager __exit__ methods

           Can also suppress exceptions the same way __exit__ methods can.
        """
        self._callbacks.append(exit)
        return exit # Allow use as a decorator

    def register(self, _cb, *args, **kwds):
        """Accepts arbitrary callbacks and arguments. Cannot suppress exceptions."""
        def _wrapper(exc_type, exc, tb):
            _cb(*args, **kwds)
        return self.register_exit(_wrapper)

    def enter_context(self, cm):
        """Accepts and automatically enters other context managers"""
        # 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)
        def _exit_wrapper(*exc_details):
            return _exit(cm, *exc_details)
        self.register_exit(_exit_wrapper)
        return result

    def close(self):
        self.__exit__(None, None, None)

    def __enter__(self):
        return self

    def __exit__(self, *exc_details):
        if not self._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._callbacks.popleft()
            if not self._callbacks:
                # Innermost callback is invoked directly
                return cb(*exc_details)
            try:
                inner_result = _invoke_next_callback(exc_details)
            except:
                cb_result = cb(*sys.exc_info())
                # Check if this cb suppressed the inner exception
                if not cb_result:
                    raise
            else:
                # Check if inner cb suppressed the original exception
                if inner_result:
                    exc_details = (None, None, None)
                cb_result = cb(*exc_details) or inner_result
            return cb_result
        # Kick off the recursive chain
        return _invoke_next_callback(exc_details)

A toy example that illustrates the LIFO execution of callbacks (necessary since callbacks registered later may depend on earlier resources):

>>> with CleanupManager() as cm:
...    @cm.register  # Registered first => cleaned up last
...    def _exit():
...        print ("Hello world!")
...    class TestCM(object):
...        def __enter__(self):
...            print ("Entered!")
...        def __exit__(self, *exc_details):
...            print ("Exited!")
...    cm.enter_context(TestCM())
... 
Entered!
Exited!
Hello world!

Cleaning up a list of files:

with contextlib.CleanupManager() as cm:
    files = [cm.enter_context(open(fname)) for fname in names]
    # All files will be closed when we leave the context

3 comments

James William Pye 12 years, 4 months ago  # | flag

Nice. Tho, I'm curious about your thoughts on handling benign [__exit__] errors on cleanup.

I think any case could be handled by CM composition of some kind, but I'm currently imagining that being a little irritating.. Perhaps better stated, I'm thinking about cases where the caller wants the innermost result in the situation of outer result exceptions. (transient error filter?)

However, that is out of scope if we are strictly thinking about nested() done right.

Nick Coghlan (author) 12 years, 4 months ago  # | flag

Context manager exit methods in general don't have any information other than whether they received an exception or not. They never know where that exception came from (i.e. the contained block or some other context manager's exit method).

Basically, there's no such thing as a benign __exit__ error: they either return a false value (saying if there was an exception, reraise it, otherwise continue as normal) or a true one (saying to continue as normal, even if there was an exception). If they throw an exception themselves, it's as if it was thrown from the body of the with statement.

CleanupManager doesn't change any of that, but it does give you a few more options for structuring your exit handlers - such as pushing a "suppress transient errors" CM as the first one on the stack, or creating two separate cleanup managers. (Hmm, I'm tempted to change the name of this recipe to ContextStack and update the API accordingly...)

James William Pye 12 years, 4 months ago  # | flag

Yeah, "ContextStack" is a better name. "CleanupManager" lends to something a bit more abstract, which is probably what got me thinking about superfluous resource finalization/release errors in the first place.