Gitiles
Code Review Sign In
gerrit-public.fairphone.software / platform / external / python / cpython3 / 088cbf2d390f6caeb021f05909e1d0b2e506b332 / . / Lib / contextlib.py
blob: 868fa6c43d77088d0f6e9fa327330f5fb8828324 [file] [log] [blame]
Guido van Rossum1a5e21e2006-02-28 21:57:43 +0000[diff] [blame]1"""Utilities for with-statement contexts. See PEP 343."""
2
3import sys
Nick Coghlan3267a302012-05-21 22:54:43 +1000[diff] [blame]4from collections import deque
Christian Heimes81ee3ef2008-05-04 22:42:01 +0000[diff] [blame]5from functools import wraps
Guido van Rossum1a5e21e2006-02-28 21:57:43 +0000[diff] [blame]6
Raymond Hettinger088cbf22013-10-10 00:46:57 -0700[diff] [blame^]7__all__ = ["contextmanager", "closing", "ContextDecorator", "ExitStack",
8 "ignored", "redirect_stdout"]
Guido van Rossum1a5e21e2006-02-28 21:57:43 +0000[diff] [blame]9
Michael Foordb3a89842010-06-30 12:17:50 +0000[diff] [blame]10
11class ContextDecorator(object):
12 "A base class or mixin that enables context managers to work as decorators."
Nick Coghlan0ded3e32011-05-05 23:49:25 +1000[diff] [blame]13
14 def _recreate_cm(self):
15 """Return a recreated instance of self.
Nick Coghlanfdc2c552011-05-06 00:02:12 +1000[diff] [blame]16
Nick Coghlan3267a302012-05-21 22:54:43 +1000[diff] [blame]17 Allows an otherwise one-shot context manager like
Nick Coghlan0ded3e32011-05-05 23:49:25 +1000[diff] [blame]18 _GeneratorContextManager to support use as
Nick Coghlan3267a302012-05-21 22:54:43 +1000[diff] [blame]19 a decorator via implicit recreation.
Nick Coghlanfdc2c552011-05-06 00:02:12 +1000[diff] [blame]20
Nick Coghlan3267a302012-05-21 22:54:43 +1000[diff] [blame]21 This is a private interface just for _GeneratorContextManager.
22 See issue #11647 for details.
Nick Coghlan0ded3e32011-05-05 23:49:25 +1000[diff] [blame]23 """
24 return self
25
Michael Foordb3a89842010-06-30 12:17:50 +0000[diff] [blame]26 def __call__(self, func):
27 @wraps(func)
28 def inner(*args, **kwds):
Nick Coghlan0ded3e32011-05-05 23:49:25 +1000[diff] [blame]29 with self._recreate_cm():
Michael Foordb3a89842010-06-30 12:17:50 +0000[diff] [blame]30 return func(*args, **kwds)
31 return inner
32
33
Antoine Pitrou67b212e2011-01-08 09:55:31 +0000[diff] [blame]34class _GeneratorContextManager(ContextDecorator):
Guido van Rossum1a5e21e2006-02-28 21:57:43 +0000[diff] [blame]35 """Helper for @contextmanager decorator."""
36
Nick Coghlan0ded3e32011-05-05 23:49:25 +1000[diff] [blame]37 def __init__(self, func, *args, **kwds):
38 self.gen = func(*args, **kwds)
39 self.func, self.args, self.kwds = func, args, kwds
40
41 def _recreate_cm(self):
42 # _GCM instances are one-shot context managers, so the
43 # CM must be recreated each time a decorated function is
44 # called
45 return self.__class__(self.func, *self.args, **self.kwds)
Guido van Rossum1a5e21e2006-02-28 21:57:43 +0000[diff] [blame]46
Guido van Rossum1a5e21e2006-02-28 21:57:43 +0000[diff] [blame]47 def __enter__(self):
48 try:
Georg Brandla18af4e2007-04-21 15:47:16 +0000[diff] [blame]49 return next(self.gen)
Guido van Rossum1a5e21e2006-02-28 21:57:43 +0000[diff] [blame]50 except StopIteration:
51 raise RuntimeError("generator didn't yield")
52
53 def __exit__(self, type, value, traceback):
54 if type is None:
55 try:
Georg Brandla18af4e2007-04-21 15:47:16 +0000[diff] [blame]56 next(self.gen)
Guido van Rossum1a5e21e2006-02-28 21:57:43 +0000[diff] [blame]57 except StopIteration:
58 return
59 else:
60 raise RuntimeError("generator didn't stop")
61 else:
Guido van Rossum2cc30da2007-11-02 23:46:40 +0000[diff] [blame]62 if value is None:
63 # Need to force instantiation so we can reliably
64 # tell if we get the same exception back
65 value = type()
Guido van Rossum1a5e21e2006-02-28 21:57:43 +0000[diff] [blame]66 try:
67 self.gen.throw(type, value, traceback)
Thomas Wouters49fd7fa2006-04-21 10:40:58 +0000[diff] [blame]68 raise RuntimeError("generator didn't stop after throw()")
Guido van Rossumb940e112007-01-10 16:19:56 +0000[diff] [blame]69 except StopIteration as exc:
Thomas Wouters49fd7fa2006-04-21 10:40:58 +0000[diff] [blame]70 # Suppress the exception *unless* it's the same exception that
71 # was passed to throw(). This prevents a StopIteration
72 # raised inside the "with" statement from being suppressed
73 return exc is not value
74 except:
75 # only re-raise if it's *not* the exception that was
76 # passed to throw(), because __exit__() must not raise
77 # an exception unless __exit__() itself failed. But throw()
78 # has to raise the exception to signal propagation, so this
79 # fixes the impedance mismatch between the throw() protocol
80 # and the __exit__() protocol.
81 #
82 if sys.exc_info()[1] is not value:
83 raise
Guido van Rossum1a5e21e2006-02-28 21:57:43 +0000[diff] [blame]84
85
86def contextmanager(func):
87 """@contextmanager decorator.
88
89 Typical usage:
90
91 @contextmanager
92 def some_generator(<arguments>):
93 <setup>
94 try:
95 yield <value>
96 finally:
97 <cleanup>
98
99 This makes this:
100
101 with some_generator(<arguments>) as <variable>:
102 <body>
103
104 equivalent to this:
105
106 <setup>
107 try:
108 <variable> = <value>
109 <body>
110 finally:
111 <cleanup>
112
113 """
Christian Heimes81ee3ef2008-05-04 22:42:01 +0000[diff] [blame]114 @wraps(func)
Guido van Rossum1a5e21e2006-02-28 21:57:43 +0000[diff] [blame]115 def helper(*args, **kwds):
Nick Coghlan0ded3e32011-05-05 23:49:25 +1000[diff] [blame]116 return _GeneratorContextManager(func, *args, **kwds)
Guido van Rossum1a5e21e2006-02-28 21:57:43 +0000[diff] [blame]117 return helper
118
119
Thomas Wouters477c8d52006-05-27 19:21:47 +0000[diff] [blame]120class closing(object):
121 """Context to automatically close something at the end of a block.
Guido van Rossum1a5e21e2006-02-28 21:57:43 +0000[diff] [blame]122
123 Code like this:
124
125 with closing(<module>.open(<arguments>)) as f:
126 <block>
127
128 is equivalent to this:
129
130 f = <module>.open(<arguments>)
131 try:
132 <block>
133 finally:
134 f.close()
135
136 """
Thomas Wouters477c8d52006-05-27 19:21:47 +0000[diff] [blame]137 def __init__(self, thing):
138 self.thing = thing
139 def __enter__(self):
140 return self.thing
141 def __exit__(self, *exc_info):
142 self.thing.close()
Nick Coghlan3267a302012-05-21 22:54:43 +1000[diff] [blame]143
Raymond Hettinger088cbf22013-10-10 00:46:57 -0700[diff] [blame^]144class redirect_stdout:
145 """Context manager for temporarily redirecting stdout to another file
146
147 # How to send help() to stderr
148
149 with redirect_stdout(sys.stderr):
150 help(dir)
151
152 # How to write help() to a file
153
154 with open('help.txt', 'w') as f:
155 with redirect_stdout(f):
156 help(pow)
157
158 # How to capture disassembly to a string
159
160 import dis
161 import io
162
163 f = io.StringIO()
164 with redirect_stdout(f):
165 dis.dis('x**2 - y**2')
166 s = f.getvalue()
167
168 """
169
170 def __init__(self, new_target):
171 self.new_target = new_target
172
173 def __enter__(self):
174 self.old_target = sys.stdout
175 sys.stdout = self.new_target
176 return self.new_target
177
178 def __exit__(self, exctype, excinst, exctb):
179 sys.stdout = self.old_target
180
Raymond Hettingere318a882013-03-10 22:26:51 -0700[diff] [blame]181@contextmanager
182def ignored(*exceptions):
Ezio Melotti9a3777e2013-08-17 15:53:55 +0300[diff] [blame]183 """Context manager to ignore specified exceptions
Raymond Hettingere318a882013-03-10 22:26:51 -0700[diff] [blame]184
185 with ignored(OSError):
186 os.remove(somefile)
187
188 """
189 try:
190 yield
191 except exceptions:
192 pass
Nick Coghlan3267a302012-05-21 22:54:43 +1000[diff] [blame]193
194# Inspired by discussions on http://bugs.python.org/issue13585
195class ExitStack(object):
196 """Context manager for dynamic management of a stack of exit callbacks
197
198 For example:
199
200 with ExitStack() as stack:
201 files = [stack.enter_context(open(fname)) for fname in filenames]
202 # All opened files will automatically be closed at the end of
203 # the with statement, even if attempts to open files later
Andrew Svetlov5b898402012-12-18 21:26:36 +0200[diff] [blame]204 # in the list raise an exception
Nick Coghlan3267a302012-05-21 22:54:43 +1000[diff] [blame]205
206 """
207 def __init__(self):
208 self._exit_callbacks = deque()
209
210 def pop_all(self):
211 """Preserve the context stack by transferring it to a new instance"""
212 new_stack = type(self)()
213 new_stack._exit_callbacks = self._exit_callbacks
214 self._exit_callbacks = deque()
215 return new_stack
216
217 def _push_cm_exit(self, cm, cm_exit):
218 """Helper to correctly register callbacks to __exit__ methods"""
219 def _exit_wrapper(*exc_details):
220 return cm_exit(cm, *exc_details)
221 _exit_wrapper.__self__ = cm
222 self.push(_exit_wrapper)
223
224 def push(self, exit):
225 """Registers a callback with the standard __exit__ method signature
226
227 Can suppress exceptions the same way __exit__ methods can.
228
229 Also accepts any object with an __exit__ method (registering a call
230 to the method instead of the object itself)
231 """
232 # We use an unbound method rather than a bound method to follow
233 # the standard lookup behaviour for special methods
234 _cb_type = type(exit)
235 try:
236 exit_method = _cb_type.__exit__
237 except AttributeError:
238 # Not a context manager, so assume its a callable
239 self._exit_callbacks.append(exit)
240 else:
241 self._push_cm_exit(exit, exit_method)
242 return exit # Allow use as a decorator
243
244 def callback(self, callback, *args, **kwds):
245 """Registers an arbitrary callback and arguments.
246
247 Cannot suppress exceptions.
248 """
249 def _exit_wrapper(exc_type, exc, tb):
250 callback(*args, **kwds)
251 # We changed the signature, so using @wraps is not appropriate, but
252 # setting __wrapped__ may still help with introspection
253 _exit_wrapper.__wrapped__ = callback
254 self.push(_exit_wrapper)
255 return callback # Allow use as a decorator
256
257 def enter_context(self, cm):
258 """Enters the supplied context manager
259
260 If successful, also pushes its __exit__ method as a callback and
261 returns the result of the __enter__ method.
262 """
263 # We look up the special methods on the type to match the with statement
264 _cm_type = type(cm)
265 _exit = _cm_type.__exit__
266 result = _cm_type.__enter__(cm)
267 self._push_cm_exit(cm, _exit)
268 return result
269
270 def close(self):
271 """Immediately unwind the context stack"""
272 self.__exit__(None, None, None)
273
274 def __enter__(self):
275 return self
276
277 def __exit__(self, *exc_details):
Nick Coghlan1a33b2f2013-10-01 23:24:56 +1000[diff] [blame]278 received_exc = exc_details[0] is not None
279
Nick Coghlan77452fc2012-06-01 22:48:32 +1000[diff] [blame]280 # We manipulate the exception state so it behaves as though
281 # we were actually nesting multiple with statements
282 frame_exc = sys.exc_info()[1]
283 def _fix_exception_context(new_exc, old_exc):
284 while 1:
285 exc_context = new_exc.__context__
286 if exc_context in (None, frame_exc):
287 break
288 new_exc = exc_context
289 new_exc.__context__ = old_exc
290
Nick Coghlana5bd2a12012-06-01 00:00:38 +1000[diff] [blame]291 # Callbacks are invoked in LIFO order to match the behaviour of
292 # nested context managers
293 suppressed_exc = False
Nick Coghlan1a33b2f2013-10-01 23:24:56 +1000[diff] [blame]294 pending_raise = False
Nick Coghlana5bd2a12012-06-01 00:00:38 +1000[diff] [blame]295 while self._exit_callbacks:
296 cb = self._exit_callbacks.pop()
Nick Coghlan3267a302012-05-21 22:54:43 +1000[diff] [blame]297 try:
Nick Coghlana5bd2a12012-06-01 00:00:38 +1000[diff] [blame]298 if cb(*exc_details):
299 suppressed_exc = True
Nick Coghlan1a33b2f2013-10-01 23:24:56 +1000[diff] [blame]300 pending_raise = False
Nick Coghlan3267a302012-05-21 22:54:43 +1000[diff] [blame]301 exc_details = (None, None, None)
Nick Coghlana5bd2a12012-06-01 00:00:38 +1000[diff] [blame]302 except:
303 new_exc_details = sys.exc_info()
Nick Coghlan77452fc2012-06-01 22:48:32 +1000[diff] [blame]304 # simulate the stack of exceptions by setting the context
305 _fix_exception_context(new_exc_details[1], exc_details[1])
Nick Coghlan1a33b2f2013-10-01 23:24:56 +1000[diff] [blame]306 pending_raise = True
Nick Coghlana5bd2a12012-06-01 00:00:38 +1000[diff] [blame]307 exc_details = new_exc_details
Nick Coghlan1a33b2f2013-10-01 23:24:56 +1000[diff] [blame]308 if pending_raise:
309 try:
310 # bare "raise exc_details[1]" replaces our carefully
311 # set-up context
312 fixed_ctx = exc_details[1].__context__
313 raise exc_details[1]
314 except BaseException:
315 exc_details[1].__context__ = fixed_ctx
316 raise
317 return received_exc and suppressed_exc