blob: c9793af6ef4ae58bd5528b31dceaf2de411784c0 [file] [log] [blame]
Guido van Rossum1a5e21e2006-02-28 21:57:43 +00001"""Utilities for with-statement contexts. See PEP 343."""
2
3import sys
Georg Brandl655fc702008-04-30 21:08:42 +00004from functools import wraps
Guido van Rossum1a5e21e2006-02-28 21:57:43 +00005
Nick Coghlanafd5e632006-05-03 13:02:47 +00006__all__ = ["contextmanager", "nested", "closing"]
Guido van Rossum1a5e21e2006-02-28 21:57:43 +00007
Nick Coghlanafd5e632006-05-03 13:02:47 +00008class GeneratorContextManager(object):
9 """Helper for @contextmanager decorator."""
Guido van Rossum1a5e21e2006-02-28 21:57:43 +000010
11 def __init__(self, gen):
12 self.gen = gen
13
Guido van Rossum1a5e21e2006-02-28 21:57:43 +000014 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:
Nick Coghlan3814a912007-11-02 10:09:12 +000029 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()
Guido van Rossum1a5e21e2006-02-28 21:57:43 +000033 try:
34 self.gen.throw(type, value, traceback)
Phillip J. Eby6edd2582006-03-25 00:28:24 +000035 raise RuntimeError("generator didn't stop after throw()")
Phillip J. Eby93149d92006-04-10 17:56:29 +000036 except StopIteration, exc:
Phillip J. Eby93880202006-04-03 21:20:07 +000037 # 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
Phillip J. Eby93149d92006-04-10 17:56:29 +000040 return exc is not value
Phillip J. Eby6edd2582006-03-25 00:28:24 +000041 except:
Phillip J. Ebyccc7bb42006-03-25 04:32:12 +000042 # 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 #
Phillip J. Eby6edd2582006-03-25 00:28:24 +000049 if sys.exc_info()[1] is not value:
50 raise
Guido van Rossum1a5e21e2006-02-28 21:57:43 +000051
52
Nick Coghlanafd5e632006-05-03 13:02:47 +000053def contextmanager(func):
54 """@contextmanager decorator.
Guido van Rossum1a5e21e2006-02-28 21:57:43 +000055
56 Typical usage:
57
Nick Coghlanafd5e632006-05-03 13:02:47 +000058 @contextmanager
Guido van Rossum1a5e21e2006-02-28 21:57:43 +000059 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 """
Georg Brandl655fc702008-04-30 21:08:42 +000081 @wraps(func)
Guido van Rossum1a5e21e2006-02-28 21:57:43 +000082 def helper(*args, **kwds):
Nick Coghlanafd5e632006-05-03 13:02:47 +000083 return GeneratorContextManager(func(*args, **kwds))
Guido van Rossum1a5e21e2006-02-28 21:57:43 +000084 return helper
85
86
Nick Coghlanafd5e632006-05-03 13:02:47 +000087@contextmanager
Guido van Rossumda5b7012006-05-02 19:47:52 +000088def nested(*managers):
Guido van Rossum1a5e21e2006-02-28 21:57:43 +000089 """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 = []
Guido van Rossumf6694362006-03-10 02:28:35 +0000106 exc = (None, None, None)
Guido van Rossum1a5e21e2006-02-28 21:57:43 +0000107 try:
Georg Brandlac4018a2007-08-23 18:11:33 +0000108 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()
Guido van Rossum1a5e21e2006-02-28 21:57:43 +0000116 finally:
117 while exits:
118 exit = exits.pop()
119 try:
Guido van Rossumf6694362006-03-10 02:28:35 +0000120 if exit(*exc):
121 exc = (None, None, None)
Guido van Rossum1a5e21e2006-02-28 21:57:43 +0000122 except:
123 exc = sys.exc_info()
Guido van Rossum1a5e21e2006-02-28 21:57:43 +0000124 if exc != (None, None, None):
Nick Coghlanda2268f2006-04-24 04:37:15 +0000125 # 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]
Guido van Rossum1a5e21e2006-02-28 21:57:43 +0000129
130
Nick Coghlana7e820a2006-04-25 10:56:51 +0000131class closing(object):
132 """Context to automatically close something at the end of a block.
Guido van Rossum1a5e21e2006-02-28 21:57:43 +0000133
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 """
Nick Coghlana7e820a2006-04-25 10:56:51 +0000148 def __init__(self, thing):
149 self.thing = thing
Nick Coghlana7e820a2006-04-25 10:56:51 +0000150 def __enter__(self):
151 return self.thing
152 def __exit__(self, *exc_info):
153 self.thing.close()