1"""Utilities for with-statement contexts.  See PEP 343."""
2
3import sys
4from functools import wraps
5from warnings import warn
6
7__all__ = ["contextmanager", "nested", "closing"]
8
9class GeneratorContextManager(object):
10    """Helper for @contextmanager decorator."""
11
12    def __init__(self, gen):
13        self.gen = gen
14
15    def __enter__(self):
16        try:
17            return self.gen.next()
18        except StopIteration:
19            raise RuntimeError("generator didn't yield")
20
21    def __exit__(self, type, value, traceback):
22        if type is None:
23            try:
24                self.gen.next()
25            except StopIteration:
26                return
27            else:
28                raise RuntimeError("generator didn't stop")
29        else:
30            if value is None:
31                # Need to force instantiation so we can reliably
32                # tell if we get the same exception back
33                value = type()
34            try:
35                self.gen.throw(type, value, traceback)
36                raise RuntimeError("generator didn't stop after throw()")
37            except StopIteration, exc:
38                # Suppress the exception *unless* it's the same exception that
39                # was passed to throw().  This prevents a StopIteration
40                # raised inside the "with" statement from being suppressed
41                return exc is not value
42            except:
43                # only re-raise if it's *not* the exception that was
44                # passed to throw(), because __exit__() must not raise
45                # an exception unless __exit__() itself failed.  But throw()
46                # has to raise the exception to signal propagation, so this
47                # fixes the impedance mismatch between the throw() protocol
48                # and the __exit__() protocol.
49                #
50                if sys.exc_info()[1] is not value:
51                    raise
52
53
54def contextmanager(func):
55    """@contextmanager decorator.
56
57    Typical usage:
58
59        @contextmanager
60        def some_generator(<arguments>):
61            <setup>
62            try:
63                yield <value>
64            finally:
65                <cleanup>
66
67    This makes this:
68
69        with some_generator(<arguments>) as <variable>:
70            <body>
71
72    equivalent to this:
73
74        <setup>
75        try:
76            <variable> = <value>
77            <body>
78        finally:
79            <cleanup>
80
81    """
82    @wraps(func)
83    def helper(*args, **kwds):
84        return GeneratorContextManager(func(*args, **kwds))
85    return helper
86
87
88@contextmanager
89def nested(*managers):
90    """Combine multiple context managers into a single nested context manager.
91
92   This function has been deprecated in favour of the multiple manager form
93   of the with statement.
94
95   The one advantage of this function over the multiple manager form of the
96   with statement is that argument unpacking allows it to be
97   used with a variable number of context managers as follows:
98
99      with nested(*managers):
100          do_something()
101
102    """
103    warn("With-statements now directly support multiple context managers",
104         DeprecationWarning, 3)
105    exits = []
106    vars = []
107    exc = (None, None, None)
108    try:
109        for mgr in managers:
110            exit = mgr.__exit__
111            enter = mgr.__enter__
112            vars.append(enter())
113            exits.append(exit)
114        yield vars
115    except:
116        exc = sys.exc_info()
117    finally:
118        while exits:
119            exit = exits.pop()
120            try:
121                if exit(*exc):
122                    exc = (None, None, None)
123            except:
124                exc = sys.exc_info()
125        if exc != (None, None, None):
126            # Don't rely on sys.exc_info() still containing
127            # the right information. Another exception may
128            # have been raised and caught by an exit method
129            raise exc[0], exc[1], exc[2]
130
131
132class closing(object):
133    """Context to automatically close something at the end of a block.
134
135    Code like this:
136
137        with closing(<module>.open(<arguments>)) as f:
138            <block>
139
140    is equivalent to this:
141
142        f = <module>.open(<arguments>)
143        try:
144            <block>
145        finally:
146            f.close()
147
148    """
149    def __init__(self, thing):
150        self.thing = thing
151    def __enter__(self):
152        return self.thing
153    def __exit__(self, *exc_info):
154        self.thing.close()
155