Gitiles
Code Review Sign In
gerrit-public.fairphone.software / platform / external / python / cpython3 / 7a34f02f9e651e8d021ceb759af102d9043b5a5b / . / Lib / contextlib.py
blob: 0b6bf71b08ca08daa5c7763ce2cd0b34a94d5518 [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
Nick Coghlan3267a302012-05-21 22:54:43 +1000[diff] [blame]7__all__ = ["contextmanager", "closing", "ContextDecorator", "ExitStack"]
Guido van Rossum1a5e21e2006-02-28 21:57:43 +0000[diff] [blame]8
Michael Foordb3a89842010-06-30 12:17:50 +0000[diff] [blame]9
10class ContextDecorator(object):
11 "A base class or mixin that enables context managers to work as decorators."
Nick Coghlan0ded3e32011-05-05 23:49:25 +1000[diff] [blame]12
13 def _recreate_cm(self):
14 """Return a recreated instance of self.
Nick Coghlanfdc2c552011-05-06 00:02:12 +1000[diff] [blame]15
Nick Coghlan3267a302012-05-21 22:54:43 +1000[diff] [blame]16 Allows an otherwise one-shot context manager like
Nick Coghlan0ded3e32011-05-05 23:49:25 +1000[diff] [blame]17 _GeneratorContextManager to support use as
Nick Coghlan3267a302012-05-21 22:54:43 +1000[diff] [blame]18 a decorator via implicit recreation.
Nick Coghlanfdc2c552011-05-06 00:02:12 +1000[diff] [blame]19
Nick Coghlan3267a302012-05-21 22:54:43 +1000[diff] [blame]20 This is a private interface just for _GeneratorContextManager.
21 See issue #11647 for details.
Nick Coghlan0ded3e32011-05-05 23:49:25 +1000[diff] [blame]22 """
23 return self
24
Michael Foordb3a89842010-06-30 12:17:50 +0000[diff] [blame]25 def __call__(self, func):
26 @wraps(func)
27 def inner(*args, **kwds):
Nick Coghlan0ded3e32011-05-05 23:49:25 +1000[diff] [blame]28 with self._recreate_cm():
Michael Foordb3a89842010-06-30 12:17:50 +0000[diff] [blame]29 return func(*args, **kwds)
30 return inner
31
32
Antoine Pitrou67b212e2011-01-08 09:55:31 +0000[diff] [blame]33class _GeneratorContextManager(ContextDecorator):
Guido van Rossum1a5e21e2006-02-28 21:57:43 +0000[diff] [blame]34 """Helper for @contextmanager decorator."""
35
Nick Coghlan0ded3e32011-05-05 23:49:25 +1000[diff] [blame]36 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 +0000[diff] [blame]45
Guido van Rossum1a5e21e2006-02-28 21:57:43 +0000[diff] [blame]46 def __enter__(self):
47 try:
Georg Brandla18af4e2007-04-21 15:47:16 +0000[diff] [blame]48 return next(self.gen)
Guido van Rossum1a5e21e2006-02-28 21:57:43 +0000[diff] [blame]49 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 +0000[diff] [blame]55 next(self.gen)
Guido van Rossum1a5e21e2006-02-28 21:57:43 +0000[diff] [blame]56 except StopIteration:
57 return
58 else:
59 raise RuntimeError("generator didn't stop")
60 else:
Guido van Rossum2cc30da2007-11-02 23:46:40 +0000[diff] [blame]61 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 +0000[diff] [blame]65 try:
66 self.gen.throw(type, value, traceback)
Thomas Wouters49fd7fa2006-04-21 10:40:58 +0000[diff] [blame]67 raise RuntimeError("generator didn't stop after throw()")
Guido van Rossumb940e112007-01-10 16:19:56 +0000[diff] [blame]68 except StopIteration as exc:
Thomas Wouters49fd7fa2006-04-21 10:40:58 +0000[diff] [blame]69 # 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 +0000[diff] [blame]83
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 +0000[diff] [blame]113 @wraps(func)
Guido van Rossum1a5e21e2006-02-28 21:57:43 +0000[diff] [blame]114 def helper(*args, **kwds):
Nick Coghlan0ded3e32011-05-05 23:49:25 +1000[diff] [blame]115 return _GeneratorContextManager(func, *args, **kwds)
Guido van Rossum1a5e21e2006-02-28 21:57:43 +0000[diff] [blame]116 return helper
117
118
Thomas Wouters477c8d52006-05-27 19:21:47 +0000[diff] [blame]119class closing(object):
120 """Context to automatically close something at the end of a block.
Guido van Rossum1a5e21e2006-02-28 21:57:43 +0000[diff] [blame]121
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 +0000[diff] [blame]136 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 +1000[diff] [blame]142
143
144# Inspired by discussions on http://bugs.python.org/issue13585
145class ExitStack(object):
146 """Context manager for dynamic management of a stack of exit callbacks
147
148 For example:
149
150 with ExitStack() as stack:
151 files = [stack.enter_context(open(fname)) for fname in filenames]
152 # All opened files will automatically be closed at the end of
153 # the with statement, even if attempts to open files later
Andrew Svetlov5b898402012-12-18 21:26:36 +0200[diff] [blame]154 # in the list raise an exception
Nick Coghlan3267a302012-05-21 22:54:43 +1000[diff] [blame]155
156 """
157 def __init__(self):
158 self._exit_callbacks = deque()
159
160 def pop_all(self):
161 """Preserve the context stack by transferring it to a new instance"""
162 new_stack = type(self)()
163 new_stack._exit_callbacks = self._exit_callbacks
164 self._exit_callbacks = deque()
165 return new_stack
166
167 def _push_cm_exit(self, cm, cm_exit):
168 """Helper to correctly register callbacks to __exit__ methods"""
169 def _exit_wrapper(*exc_details):
170 return cm_exit(cm, *exc_details)
171 _exit_wrapper.__self__ = cm
172 self.push(_exit_wrapper)
173
174 def push(self, exit):
175 """Registers a callback with the standard __exit__ method signature
176
177 Can suppress exceptions the same way __exit__ methods can.
178
179 Also accepts any object with an __exit__ method (registering a call
180 to the method instead of the object itself)
181 """
182 # We use an unbound method rather than a bound method to follow
183 # the standard lookup behaviour for special methods
184 _cb_type = type(exit)
185 try:
186 exit_method = _cb_type.__exit__
187 except AttributeError:
188 # Not a context manager, so assume its a callable
189 self._exit_callbacks.append(exit)
190 else:
191 self._push_cm_exit(exit, exit_method)
192 return exit # Allow use as a decorator
193
194 def callback(self, callback, *args, **kwds):
195 """Registers an arbitrary callback and arguments.
196
197 Cannot suppress exceptions.
198 """
199 def _exit_wrapper(exc_type, exc, tb):
200 callback(*args, **kwds)
201 # We changed the signature, so using @wraps is not appropriate, but
202 # setting __wrapped__ may still help with introspection
203 _exit_wrapper.__wrapped__ = callback
204 self.push(_exit_wrapper)
205 return callback # Allow use as a decorator
206
207 def enter_context(self, cm):
208 """Enters the supplied context manager
209
210 If successful, also pushes its __exit__ method as a callback and
211 returns the result of the __enter__ method.
212 """
213 # We look up the special methods on the type to match the with statement
214 _cm_type = type(cm)
215 _exit = _cm_type.__exit__
216 result = _cm_type.__enter__(cm)
217 self._push_cm_exit(cm, _exit)
218 return result
219
220 def close(self):
221 """Immediately unwind the context stack"""
222 self.__exit__(None, None, None)
223
224 def __enter__(self):
225 return self
226
227 def __exit__(self, *exc_details):
Nick Coghlan77452fc2012-06-01 22:48:32 +1000[diff] [blame]228 # We manipulate the exception state so it behaves as though
229 # we were actually nesting multiple with statements
230 frame_exc = sys.exc_info()[1]
231 def _fix_exception_context(new_exc, old_exc):
232 while 1:
233 exc_context = new_exc.__context__
234 if exc_context in (None, frame_exc):
235 break
236 new_exc = exc_context
237 new_exc.__context__ = old_exc
238
Nick Coghlana5bd2a12012-06-01 00:00:38 +1000[diff] [blame]239 # Callbacks are invoked in LIFO order to match the behaviour of
240 # nested context managers
241 suppressed_exc = False
242 while self._exit_callbacks:
243 cb = self._exit_callbacks.pop()
Nick Coghlan3267a302012-05-21 22:54:43 +1000[diff] [blame]244 try:
Nick Coghlana5bd2a12012-06-01 00:00:38 +1000[diff] [blame]245 if cb(*exc_details):
246 suppressed_exc = True
Nick Coghlan3267a302012-05-21 22:54:43 +1000[diff] [blame]247 exc_details = (None, None, None)
Nick Coghlana5bd2a12012-06-01 00:00:38 +1000[diff] [blame]248 except:
249 new_exc_details = sys.exc_info()
Nick Coghlan77452fc2012-06-01 22:48:32 +1000[diff] [blame]250 # simulate the stack of exceptions by setting the context
251 _fix_exception_context(new_exc_details[1], exc_details[1])
Nick Coghlana5bd2a12012-06-01 00:00:38 +1000[diff] [blame]252 if not self._exit_callbacks:
253 raise
254 exc_details = new_exc_details
255 return suppressed_exc