blob: aeec40e1e32685ec0f2a3266d3e29670adbe3b93 [file] [log] [blame]
Guido van Rossum1a5e21e2006-02-28 21:57:43 +00001"""Utilities for with-statement contexts. See PEP 343."""
2
3import sys
4
5__all__ = ["contextmanager", "nested", "closing"]
6
7class GeneratorContextManager(object):
8 """Helper for @contextmanager decorator."""
9
10 def __init__(self, gen):
11 self.gen = gen
12
Guido van Rossum1a5e21e2006-02-28 21:57:43 +000013 def __enter__(self):
14 try:
Georg Brandla18af4e2007-04-21 15:47:16 +000015 return next(self.gen)
Guido van Rossum1a5e21e2006-02-28 21:57:43 +000016 except StopIteration:
17 raise RuntimeError("generator didn't yield")
18
19 def __exit__(self, type, value, traceback):
20 if type is None:
21 try:
Georg Brandla18af4e2007-04-21 15:47:16 +000022 next(self.gen)
Guido van Rossum1a5e21e2006-02-28 21:57:43 +000023 except StopIteration:
24 return
25 else:
26 raise RuntimeError("generator didn't stop")
27 else:
Guido van Rossum2cc30da2007-11-02 23:46:40 +000028 if value is None:
29 # Need to force instantiation so we can reliably
30 # tell if we get the same exception back
31 value = type()
Guido van Rossum1a5e21e2006-02-28 21:57:43 +000032 try:
33 self.gen.throw(type, value, traceback)
Thomas Wouters49fd7fa2006-04-21 10:40:58 +000034 raise RuntimeError("generator didn't stop after throw()")
Guido van Rossumb940e112007-01-10 16:19:56 +000035 except StopIteration as exc:
Thomas Wouters49fd7fa2006-04-21 10:40:58 +000036 # Suppress the exception *unless* it's the same exception that
37 # was passed to throw(). This prevents a StopIteration
38 # raised inside the "with" statement from being suppressed
39 return exc is not value
40 except:
41 # only re-raise if it's *not* the exception that was
42 # passed to throw(), because __exit__() must not raise
43 # an exception unless __exit__() itself failed. But throw()
44 # has to raise the exception to signal propagation, so this
45 # fixes the impedance mismatch between the throw() protocol
46 # and the __exit__() protocol.
47 #
48 if sys.exc_info()[1] is not value:
49 raise
Guido van Rossum1a5e21e2006-02-28 21:57:43 +000050
51
52def contextmanager(func):
53 """@contextmanager decorator.
54
55 Typical usage:
56
57 @contextmanager
58 def some_generator(<arguments>):
59 <setup>
60 try:
61 yield <value>
62 finally:
63 <cleanup>
64
65 This makes this:
66
67 with some_generator(<arguments>) as <variable>:
68 <body>
69
70 equivalent to this:
71
72 <setup>
73 try:
74 <variable> = <value>
75 <body>
76 finally:
77 <cleanup>
78
79 """
80 def helper(*args, **kwds):
81 return GeneratorContextManager(func(*args, **kwds))
82 try:
83 helper.__name__ = func.__name__
84 helper.__doc__ = func.__doc__
Thomas Wouters49fd7fa2006-04-21 10:40:58 +000085 helper.__dict__ = func.__dict__
Guido van Rossum1a5e21e2006-02-28 21:57:43 +000086 except:
87 pass
88 return helper
89
90
91@contextmanager
Thomas Wouters477c8d52006-05-27 19:21:47 +000092def nested(*managers):
Guido van Rossum1a5e21e2006-02-28 21:57:43 +000093 """Support multiple context managers in a single with-statement.
94
95 Code like this:
96
97 with nested(A, B, C) as (X, Y, Z):
98 <body>
99
100 is equivalent to this:
101
102 with A as X:
103 with B as Y:
104 with C as Z:
105 <body>
106
107 """
108 exits = []
109 vars = []
Guido van Rossumf6694362006-03-10 02:28:35 +0000110 exc = (None, None, None)
Guido van Rossum1a5e21e2006-02-28 21:57:43 +0000111 try:
Guido van Rossum04110fb2007-08-24 16:32:05 +0000112 for mgr in managers:
113 exit = mgr.__exit__
114 enter = mgr.__enter__
115 vars.append(enter())
116 exits.append(exit)
117 yield vars
118 except:
119 exc = sys.exc_info()
Guido van Rossum1a5e21e2006-02-28 21:57:43 +0000120 finally:
121 while exits:
122 exit = exits.pop()
123 try:
Guido van Rossumf6694362006-03-10 02:28:35 +0000124 if exit(*exc):
125 exc = (None, None, None)
Guido van Rossum1a5e21e2006-02-28 21:57:43 +0000126 except:
127 exc = sys.exc_info()
Guido van Rossum1a5e21e2006-02-28 21:57:43 +0000128 if exc != (None, None, None):
Thomas Wouters477c8d52006-05-27 19:21:47 +0000129 # Don't rely on sys.exc_info() still containing
130 # the right information. Another exception may
131 # have been raised and caught by an exit method
Collin Winterc59dacd2007-09-01 20:27:58 +0000132 # exc[1] already has the __traceback__ attribute populated
133 raise exc[1]
Guido van Rossum1a5e21e2006-02-28 21:57:43 +0000134
135
Thomas Wouters477c8d52006-05-27 19:21:47 +0000136class closing(object):
137 """Context to automatically close something at the end of a block.
Guido van Rossum1a5e21e2006-02-28 21:57:43 +0000138
139 Code like this:
140
141 with closing(<module>.open(<arguments>)) as f:
142 <block>
143
144 is equivalent to this:
145
146 f = <module>.open(<arguments>)
147 try:
148 <block>
149 finally:
150 f.close()
151
152 """
Thomas Wouters477c8d52006-05-27 19:21:47 +0000153 def __init__(self, thing):
154 self.thing = thing
155 def __enter__(self):
156 return self.thing
157 def __exit__(self, *exc_info):
158 self.thing.close()