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