blob: 03c56da39758e0b18989b875a6e72e9da9b43474 [file] [log] [blame]
Guido van Rossum1a5e21e2006-02-28 21:57:43 +00001"""Utilities for with-statement contexts. See PEP 343."""
2
3import sys
Nick Coghlan3267a302012-05-21 22:54:43 +10004from collections import deque
Christian Heimes81ee3ef2008-05-04 22:42:01 +00005from functools import wraps
Guido van Rossum1a5e21e2006-02-28 21:57:43 +00006
Raymond Hettingere318a882013-03-10 22:26:51 -07007__all__ = ["contextmanager", "closing", "ContextDecorator", "ExitStack", "ignored"]
Guido van Rossum1a5e21e2006-02-28 21:57:43 +00008
Michael Foordb3a89842010-06-30 12:17:50 +00009
10class ContextDecorator(object):
11 "A base class or mixin that enables context managers to work as decorators."
Nick Coghlan0ded3e32011-05-05 23:49:25 +100012
13 def _recreate_cm(self):
14 """Return a recreated instance of self.
Nick Coghlanfdc2c552011-05-06 00:02:12 +100015
Nick Coghlan3267a302012-05-21 22:54:43 +100016 Allows an otherwise one-shot context manager like
Nick Coghlan0ded3e32011-05-05 23:49:25 +100017 _GeneratorContextManager to support use as
Nick Coghlan3267a302012-05-21 22:54:43 +100018 a decorator via implicit recreation.
Nick Coghlanfdc2c552011-05-06 00:02:12 +100019
Nick Coghlan3267a302012-05-21 22:54:43 +100020 This is a private interface just for _GeneratorContextManager.
21 See issue #11647 for details.
Nick Coghlan0ded3e32011-05-05 23:49:25 +100022 """
23 return self
24
Michael Foordb3a89842010-06-30 12:17:50 +000025 def __call__(self, func):
26 @wraps(func)
27 def inner(*args, **kwds):
Nick Coghlan0ded3e32011-05-05 23:49:25 +100028 with self._recreate_cm():
Michael Foordb3a89842010-06-30 12:17:50 +000029 return func(*args, **kwds)
30 return inner
31
32
Antoine Pitrou67b212e2011-01-08 09:55:31 +000033class _GeneratorContextManager(ContextDecorator):
Guido van Rossum1a5e21e2006-02-28 21:57:43 +000034 """Helper for @contextmanager decorator."""
35
Nick Coghlan0ded3e32011-05-05 23:49:25 +100036 def __init__(self, func, *args, **kwds):
37 self.gen = func(*args, **kwds)
38 self.func, self.args, self.kwds = func, args, kwds
39
40 def _recreate_cm(self):
41 # _GCM instances are one-shot context managers, so the
42 # CM must be recreated each time a decorated function is
43 # called
44 return self.__class__(self.func, *self.args, **self.kwds)
Guido van Rossum1a5e21e2006-02-28 21:57:43 +000045
Guido van Rossum1a5e21e2006-02-28 21:57:43 +000046 def __enter__(self):
47 try:
Georg Brandla18af4e2007-04-21 15:47:16 +000048 return next(self.gen)
Guido van Rossum1a5e21e2006-02-28 21:57:43 +000049 except StopIteration:
50 raise RuntimeError("generator didn't yield")
51
52 def __exit__(self, type, value, traceback):
53 if type is None:
54 try:
Georg Brandla18af4e2007-04-21 15:47:16 +000055 next(self.gen)
Guido van Rossum1a5e21e2006-02-28 21:57:43 +000056 except StopIteration:
57 return
58 else:
59 raise RuntimeError("generator didn't stop")
60 else:
Guido van Rossum2cc30da2007-11-02 23:46:40 +000061 if value is None:
62 # Need to force instantiation so we can reliably
63 # tell if we get the same exception back
64 value = type()
Guido van Rossum1a5e21e2006-02-28 21:57:43 +000065 try:
66 self.gen.throw(type, value, traceback)
Thomas Wouters49fd7fa2006-04-21 10:40:58 +000067 raise RuntimeError("generator didn't stop after throw()")
Guido van Rossumb940e112007-01-10 16:19:56 +000068 except StopIteration as exc:
Thomas Wouters49fd7fa2006-04-21 10:40:58 +000069 # Suppress the exception *unless* it's the same exception that
70 # was passed to throw(). This prevents a StopIteration
71 # raised inside the "with" statement from being suppressed
72 return exc is not value
73 except:
74 # only re-raise if it's *not* the exception that was
75 # passed to throw(), because __exit__() must not raise
76 # an exception unless __exit__() itself failed. But throw()
77 # has to raise the exception to signal propagation, so this
78 # fixes the impedance mismatch between the throw() protocol
79 # and the __exit__() protocol.
80 #
81 if sys.exc_info()[1] is not value:
82 raise
Guido van Rossum1a5e21e2006-02-28 21:57:43 +000083
84
85def contextmanager(func):
86 """@contextmanager decorator.
87
88 Typical usage:
89
90 @contextmanager
91 def some_generator(<arguments>):
92 <setup>
93 try:
94 yield <value>
95 finally:
96 <cleanup>
97
98 This makes this:
99
100 with some_generator(<arguments>) as <variable>:
101 <body>
102
103 equivalent to this:
104
105 <setup>
106 try:
107 <variable> = <value>
108 <body>
109 finally:
110 <cleanup>
111
112 """
Christian Heimes81ee3ef2008-05-04 22:42:01 +0000113 @wraps(func)
Guido van Rossum1a5e21e2006-02-28 21:57:43 +0000114 def helper(*args, **kwds):
Nick Coghlan0ded3e32011-05-05 23:49:25 +1000115 return _GeneratorContextManager(func, *args, **kwds)
Guido van Rossum1a5e21e2006-02-28 21:57:43 +0000116 return helper
117
118
Thomas Wouters477c8d52006-05-27 19:21:47 +0000119class closing(object):
120 """Context to automatically close something at the end of a block.
Guido van Rossum1a5e21e2006-02-28 21:57:43 +0000121
122 Code like this:
123
124 with closing(<module>.open(<arguments>)) as f:
125 <block>
126
127 is equivalent to this:
128
129 f = <module>.open(<arguments>)
130 try:
131 <block>
132 finally:
133 f.close()
134
135 """
Thomas Wouters477c8d52006-05-27 19:21:47 +0000136 def __init__(self, thing):
137 self.thing = thing
138 def __enter__(self):
139 return self.thing
140 def __exit__(self, *exc_info):
141 self.thing.close()
Nick Coghlan3267a302012-05-21 22:54:43 +1000142
Raymond Hettingere318a882013-03-10 22:26:51 -0700143@contextmanager
144def ignored(*exceptions):
145 """Context manager to ignore specifed exceptions
146
147 with ignored(OSError):
148 os.remove(somefile)
149
150 """
151 try:
152 yield
153 except exceptions:
154 pass
Nick Coghlan3267a302012-05-21 22:54:43 +1000155
156# Inspired by discussions on http://bugs.python.org/issue13585
157class ExitStack(object):
158 """Context manager for dynamic management of a stack of exit callbacks
159
160 For example:
161
162 with ExitStack() as stack:
163 files = [stack.enter_context(open(fname)) for fname in filenames]
164 # All opened files will automatically be closed at the end of
165 # the with statement, even if attempts to open files later
Andrew Svetlov5b898402012-12-18 21:26:36 +0200166 # in the list raise an exception
Nick Coghlan3267a302012-05-21 22:54:43 +1000167
168 """
169 def __init__(self):
170 self._exit_callbacks = deque()
171
172 def pop_all(self):
173 """Preserve the context stack by transferring it to a new instance"""
174 new_stack = type(self)()
175 new_stack._exit_callbacks = self._exit_callbacks
176 self._exit_callbacks = deque()
177 return new_stack
178
179 def _push_cm_exit(self, cm, cm_exit):
180 """Helper to correctly register callbacks to __exit__ methods"""
181 def _exit_wrapper(*exc_details):
182 return cm_exit(cm, *exc_details)
183 _exit_wrapper.__self__ = cm
184 self.push(_exit_wrapper)
185
186 def push(self, exit):
187 """Registers a callback with the standard __exit__ method signature
188
189 Can suppress exceptions the same way __exit__ methods can.
190
191 Also accepts any object with an __exit__ method (registering a call
192 to the method instead of the object itself)
193 """
194 # We use an unbound method rather than a bound method to follow
195 # the standard lookup behaviour for special methods
196 _cb_type = type(exit)
197 try:
198 exit_method = _cb_type.__exit__
199 except AttributeError:
200 # Not a context manager, so assume its a callable
201 self._exit_callbacks.append(exit)
202 else:
203 self._push_cm_exit(exit, exit_method)
204 return exit # Allow use as a decorator
205
206 def callback(self, callback, *args, **kwds):
207 """Registers an arbitrary callback and arguments.
208
209 Cannot suppress exceptions.
210 """
211 def _exit_wrapper(exc_type, exc, tb):
212 callback(*args, **kwds)
213 # We changed the signature, so using @wraps is not appropriate, but
214 # setting __wrapped__ may still help with introspection
215 _exit_wrapper.__wrapped__ = callback
216 self.push(_exit_wrapper)
217 return callback # Allow use as a decorator
218
219 def enter_context(self, cm):
220 """Enters the supplied context manager
221
222 If successful, also pushes its __exit__ method as a callback and
223 returns the result of the __enter__ method.
224 """
225 # We look up the special methods on the type to match the with statement
226 _cm_type = type(cm)
227 _exit = _cm_type.__exit__
228 result = _cm_type.__enter__(cm)
229 self._push_cm_exit(cm, _exit)
230 return result
231
232 def close(self):
233 """Immediately unwind the context stack"""
234 self.__exit__(None, None, None)
235
236 def __enter__(self):
237 return self
238
239 def __exit__(self, *exc_details):
Nick Coghlan77452fc2012-06-01 22:48:32 +1000240 # We manipulate the exception state so it behaves as though
241 # we were actually nesting multiple with statements
242 frame_exc = sys.exc_info()[1]
243 def _fix_exception_context(new_exc, old_exc):
244 while 1:
245 exc_context = new_exc.__context__
246 if exc_context in (None, frame_exc):
247 break
248 new_exc = exc_context
249 new_exc.__context__ = old_exc
250
Nick Coghlana5bd2a12012-06-01 00:00:38 +1000251 # Callbacks are invoked in LIFO order to match the behaviour of
252 # nested context managers
253 suppressed_exc = False
254 while self._exit_callbacks:
255 cb = self._exit_callbacks.pop()
Nick Coghlan3267a302012-05-21 22:54:43 +1000256 try:
Nick Coghlana5bd2a12012-06-01 00:00:38 +1000257 if cb(*exc_details):
258 suppressed_exc = True
Nick Coghlan3267a302012-05-21 22:54:43 +1000259 exc_details = (None, None, None)
Nick Coghlana5bd2a12012-06-01 00:00:38 +1000260 except:
261 new_exc_details = sys.exc_info()
Nick Coghlan77452fc2012-06-01 22:48:32 +1000262 # simulate the stack of exceptions by setting the context
263 _fix_exception_context(new_exc_details[1], exc_details[1])
Nick Coghlana5bd2a12012-06-01 00:00:38 +1000264 if not self._exit_callbacks:
265 raise
266 exc_details = new_exc_details
267 return suppressed_exc