Python code coverage for Lib/contextlib.py

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