1
+ """Utilities for with-statement contexts. See PEP 343."""
2
+
3
+ import sys
4
+ from collections import deque
5
+ from functools import wraps
6
+
7
+ __all__ = ["contextmanager" , "closing" , "ContextDecorator" , "ExitStack" ,
8
+ "redirect_stdout" , "suppress" ]
9
+
10
+
11
+ class ContextDecorator (object ):
12
+ "A base class or mixin that enables context managers to work as decorators."
13
+
14
+ def _recreate_cm (self ):
15
+ """Return a recreated instance of self.
16
+
17
+ Allows an otherwise one-shot context manager like
18
+ _GeneratorContextManager to support use as
19
+ a decorator via implicit recreation.
20
+
21
+ This is a private interface just for _GeneratorContextManager.
22
+ See issue #11647 for details.
23
+ """
24
+ return self
25
+
26
+ def __call__ (self , func ):
27
+ @wraps (func )
28
+ def inner (* args , ** kwds ):
29
+ with self ._recreate_cm ():
30
+ return func (* args , ** kwds )
31
+ return inner
32
+
33
+
34
+ class _GeneratorContextManager (ContextDecorator ):
35
+ """Helper for @contextmanager decorator."""
36
+
37
+ def __init__ (self , func , * args , ** kwds ):
38
+ self .gen = func (* args , ** kwds )
39
+ self .func , self .args , self .kwds = func , args , kwds
40
+ # Issue 19330: ensure context manager instances have good docstrings
41
+ doc = getattr (func , "__doc__" , None )
42
+ if doc is None :
43
+ doc = type (self ).__doc__
44
+ self .__doc__ = doc
45
+ # Unfortunately, this still doesn't provide good help output when
46
+ # inspecting the created context manager instances, since pydoc
47
+ # currently bypasses the instance docstring and shows the docstring
48
+ # for the class instead.
49
+ # See http://bugs.python.org/issue19404 for more details.
50
+
51
+ def _recreate_cm (self ):
52
+ # _GCM instances are one-shot context managers, so the
53
+ # CM must be recreated each time a decorated function is
54
+ # called
55
+ return self .__class__ (self .func , * self .args , ** self .kwds )
56
+
57
+ def __enter__ (self ):
58
+ try :
59
+ return next (self .gen )
60
+ except StopIteration :
61
+ raise RuntimeError ("generator didn't yield" ) from None
62
+
63
+ def __exit__ (self , type , value , traceback ):
64
+ if type is None :
65
+ try :
66
+ next (self .gen )
67
+ except StopIteration :
68
+ return
69
+ else :
70
+ raise RuntimeError ("generator didn't stop" )
71
+ else :
72
+ if value is None :
73
+ # Need to force instantiation so we can reliably
74
+ # tell if we get the same exception back
75
+ value = type ()
76
+ try :
77
+ self .gen .throw (type , value , traceback )
78
+ raise RuntimeError ("generator didn't stop after throw()" )
79
+ except StopIteration as exc :
80
+ # Suppress the exception *unless* it's the same exception that
81
+ # was passed to throw(). This prevents a StopIteration
82
+ # raised inside the "with" statement from being suppressed
83
+ return exc is not value
84
+ except :
85
+ # only re-raise if it's *not* the exception that was
86
+ # passed to throw(), because __exit__() must not raise
87
+ # an exception unless __exit__() itself failed. But throw()
88
+ # has to raise the exception to signal propagation, so this
89
+ # fixes the impedance mismatch between the throw() protocol
90
+ # and the __exit__() protocol.
91
+ #
92
+ if sys .exc_info ()[1 ] is not value :
93
+ raise
94
+
95
+
96
+ def contextmanager (func ):
97
+ """@contextmanager decorator.
98
+
99
+ Typical usage:
100
+
101
+ @contextmanager
102
+ def some_generator(<arguments>):
103
+ <setup>
104
+ try:
105
+ yield <value>
106
+ finally:
107
+ <cleanup>
108
+
109
+ This makes this:
110
+
111
+ with some_generator(<arguments>) as <variable>:
112
+ <body>
113
+
114
+ equivalent to this:
115
+
116
+ <setup>
117
+ try:
118
+ <variable> = <value>
119
+ <body>
120
+ finally:
121
+ <cleanup>
122
+
123
+ """
124
+ @wraps (func )
125
+ def helper (* args , ** kwds ):
126
+ return _GeneratorContextManager (func , * args , ** kwds )
127
+ return helper
128
+
129
+
130
+ class closing (object ):
131
+ """Context to automatically close something at the end of a block.
132
+
133
+ Code like this:
134
+
135
+ with closing(<module>.open(<arguments>)) as f:
136
+ <block>
137
+
138
+ is equivalent to this:
139
+
140
+ f = <module>.open(<arguments>)
141
+ try:
142
+ <block>
143
+ finally:
144
+ f.close()
145
+
146
+ """
147
+ def __init__ (self , thing ):
148
+ self .thing = thing
149
+ def __enter__ (self ):
150
+ return self .thing
151
+ def __exit__ (self , * exc_info ):
152
+ self .thing .close ()
153
+
154
+ class redirect_stdout :
155
+ """Context manager for temporarily redirecting stdout to another file
156
+
157
+ # How to send help() to stderr
158
+ with redirect_stdout(sys.stderr):
159
+ help(dir)
160
+
161
+ # How to write help() to a file
162
+ with open('help.txt', 'w') as f:
163
+ with redirect_stdout(f):
164
+ help(pow)
165
+ """
166
+
167
+ def __init__ (self , new_target ):
168
+ self ._new_target = new_target
169
+ # We use a list of old targets to make this CM re-entrant
170
+ self ._old_targets = []
171
+
172
+ def __enter__ (self ):
173
+ self ._old_targets .append (sys .stdout )
174
+ sys .stdout = self ._new_target
175
+ return self ._new_target
176
+
177
+ def __exit__ (self , exctype , excinst , exctb ):
178
+ sys .stdout = self ._old_targets .pop ()
179
+
180
+
181
+ class suppress :
182
+ """Context manager to suppress specified exceptions
183
+
184
+ After the exception is suppressed, execution proceeds with the next
185
+ statement following the with statement.
186
+
187
+ with suppress(FileNotFoundError):
188
+ os.remove(somefile)
189
+ # Execution still resumes here if the file was already removed
190
+ """
191
+
192
+ def __init__ (self , * exceptions ):
193
+ self ._exceptions = exceptions
194
+
195
+ def __enter__ (self ):
196
+ pass
197
+
198
+ def __exit__ (self , exctype , excinst , exctb ):
199
+ # Unlike isinstance and issubclass, CPython exception handling
200
+ # currently only looks at the concrete type hierarchy (ignoring
201
+ # the instance and subclass checking hooks). While Guido considers
202
+ # that a bug rather than a feature, it's a fairly hard one to fix
203
+ # due to various internal implementation details. suppress provides
204
+ # the simpler issubclass based semantics, rather than trying to
205
+ # exactly reproduce the limitations of the CPython interpreter.
206
+ #
207
+ # See http://bugs.python.org/issue12029 for more details
208
+ return exctype is not None and issubclass (exctype , self ._exceptions )
209
+
210
+
211
+ # Inspired by discussions on http://bugs.python.org/issue13585
212
+ class ExitStack (object ):
213
+ """Context manager for dynamic management of a stack of exit callbacks
214
+
215
+ For example:
216
+
217
+ with ExitStack() as stack:
218
+ files = [stack.enter_context(open(fname)) for fname in filenames]
219
+ # All opened files will automatically be closed at the end of
220
+ # the with statement, even if attempts to open files later
221
+ # in the list raise an exception
222
+
223
+ """
224
+ def __init__ (self ):
225
+ self ._exit_callbacks = deque ()
226
+
227
+ def pop_all (self ):
228
+ """Preserve the context stack by transferring it to a new instance"""
229
+ new_stack = type (self )()
230
+ new_stack ._exit_callbacks = self ._exit_callbacks
231
+ self ._exit_callbacks = deque ()
232
+ return new_stack
233
+
234
+ def _push_cm_exit (self , cm , cm_exit ):
235
+ """Helper to correctly register callbacks to __exit__ methods"""
236
+ def _exit_wrapper (* exc_details ):
237
+ return cm_exit (cm , * exc_details )
238
+ _exit_wrapper .__self__ = cm
239
+ self .push (_exit_wrapper )
240
+
241
+ def push (self , exit ):
242
+ """Registers a callback with the standard __exit__ method signature
243
+
244
+ Can suppress exceptions the same way __exit__ methods can.
245
+
246
+ Also accepts any object with an __exit__ method (registering a call
247
+ to the method instead of the object itself)
248
+ """
249
+ # We use an unbound method rather than a bound method to follow
250
+ # the standard lookup behaviour for special methods
251
+ _cb_type = type (exit )
252
+ try :
253
+ exit_method = _cb_type .__exit__
254
+ except AttributeError :
255
+ # Not a context manager, so assume its a callable
256
+ self ._exit_callbacks .append (exit )
257
+ else :
258
+ self ._push_cm_exit (exit , exit_method )
259
+ return exit # Allow use as a decorator
260
+
261
+ def callback (self , callback , * args , ** kwds ):
262
+ """Registers an arbitrary callback and arguments.
263
+
264
+ Cannot suppress exceptions.
265
+ """
266
+ def _exit_wrapper (exc_type , exc , tb ):
267
+ callback (* args , ** kwds )
268
+ # We changed the signature, so using @wraps is not appropriate, but
269
+ # setting __wrapped__ may still help with introspection
270
+ _exit_wrapper .__wrapped__ = callback
271
+ self .push (_exit_wrapper )
272
+ return callback # Allow use as a decorator
273
+
274
+ def enter_context (self , cm ):
275
+ """Enters the supplied context manager
276
+
277
+ If successful, also pushes its __exit__ method as a callback and
278
+ returns the result of the __enter__ method.
279
+ """
280
+ # We look up the special methods on the type to match the with statement
281
+ _cm_type = type (cm )
282
+ _exit = _cm_type .__exit__
283
+ result = _cm_type .__enter__ (cm )
284
+ self ._push_cm_exit (cm , _exit )
285
+ return result
286
+
287
+ def close (self ):
288
+ """Immediately unwind the context stack"""
289
+ self .__exit__ (None , None , None )
290
+
291
+ def __enter__ (self ):
292
+ return self
293
+
294
+ def __exit__ (self , * exc_details ):
295
+ received_exc = exc_details [0 ] is not None
296
+
297
+ # We manipulate the exception state so it behaves as though
298
+ # we were actually nesting multiple with statements
299
+ frame_exc = sys .exc_info ()[1 ]
300
+ def _fix_exception_context (new_exc , old_exc ):
301
+ # Context may not be correct, so find the end of the chain
302
+ while 1 :
303
+ exc_context = new_exc .__context__
304
+ if exc_context is old_exc :
305
+ # Context is already set correctly (see issue 20317)
306
+ return
307
+ if exc_context is None or exc_context is frame_exc :
308
+ break
309
+ new_exc = exc_context
310
+ # Change the end of the chain to point to the exception
311
+ # we expect it to reference
312
+ new_exc .__context__ = old_exc
313
+
314
+ # Callbacks are invoked in LIFO order to match the behaviour of
315
+ # nested context managers
316
+ suppressed_exc = False
317
+ pending_raise = False
318
+ while self ._exit_callbacks :
319
+ cb = self ._exit_callbacks .pop ()
320
+ try :
321
+ if cb (* exc_details ):
322
+ suppressed_exc = True
323
+ pending_raise = False
324
+ exc_details = (None , None , None )
325
+ except :
326
+ new_exc_details = sys .exc_info ()
327
+ # simulate the stack of exceptions by setting the context
328
+ _fix_exception_context (new_exc_details [1 ], exc_details [1 ])
329
+ pending_raise = True
330
+ exc_details = new_exc_details
331
+ if pending_raise :
332
+ try :
333
+ # bare "raise exc_details[1]" replaces our carefully
334
+ # set-up context
335
+ fixed_ctx = exc_details [1 ].__context__
336
+ raise exc_details [1 ]
337
+ except BaseException :
338
+ exc_details [1 ].__context__ = fixed_ctx
339
+ raise
340
+ return received_exc and suppressed_exc
0 commit comments