blob: 50943c1a17e9cb0964926d55b59ca8ea28b6936e [file] [log] [blame]
Guido van Rossum1a5e21e2006-02-28 21:57:43 +00001"""Unit tests for contextlib.py, and other context managers."""
2
Raymond Hettinger088cbf22013-10-10 00:46:57 -07003import io
R. David Murray378c0cf2010-02-24 01:46:21 +00004import sys
Guido van Rossum1a5e21e2006-02-28 21:57:43 +00005import tempfile
Antoine Pitroua6a4dc82017-09-07 18:56:24 +02006import threading
Guido van Rossum1a5e21e2006-02-28 21:57:43 +00007import unittest
Guido van Rossum1a5e21e2006-02-28 21:57:43 +00008from contextlib import * # Tests __all__
Benjamin Petersonee8712c2008-05-20 21:35:26 +00009from test import support
Hai Shi96a6a6d2020-07-09 21:25:10 +080010from test.support import os_helper
Martin Teichmanndd0e0872018-01-28 05:17:46 +010011import weakref
Guido van Rossum1a5e21e2006-02-28 21:57:43 +000012
Florent Xicluna41fe6152010-04-02 18:52:12 +000013
Brett Cannon9e080e02016-04-08 12:15:27 -070014class TestAbstractContextManager(unittest.TestCase):
15
16 def test_enter(self):
17 class DefaultEnter(AbstractContextManager):
18 def __exit__(self, *args):
19 super().__exit__(*args)
20
21 manager = DefaultEnter()
22 self.assertIs(manager.__enter__(), manager)
23
24 def test_exit_is_abstract(self):
25 class MissingExit(AbstractContextManager):
26 pass
27
28 with self.assertRaises(TypeError):
29 MissingExit()
30
31 def test_structural_subclassing(self):
32 class ManagerFromScratch:
33 def __enter__(self):
34 return self
35 def __exit__(self, exc_type, exc_value, traceback):
36 return None
37
38 self.assertTrue(issubclass(ManagerFromScratch, AbstractContextManager))
39
40 class DefaultEnter(AbstractContextManager):
41 def __exit__(self, *args):
42 super().__exit__(*args)
43
44 self.assertTrue(issubclass(DefaultEnter, AbstractContextManager))
45
Jelle Zijlstra57161aa2017-06-09 08:21:47 -070046 class NoEnter(ManagerFromScratch):
47 __enter__ = None
48
49 self.assertFalse(issubclass(NoEnter, AbstractContextManager))
50
51 class NoExit(ManagerFromScratch):
52 __exit__ = None
53
54 self.assertFalse(issubclass(NoExit, AbstractContextManager))
55
Brett Cannon9e080e02016-04-08 12:15:27 -070056
Guido van Rossum1a5e21e2006-02-28 21:57:43 +000057class ContextManagerTestCase(unittest.TestCase):
58
59 def test_contextmanager_plain(self):
60 state = []
61 @contextmanager
62 def woohoo():
63 state.append(1)
64 yield 42
65 state.append(999)
66 with woohoo() as x:
67 self.assertEqual(state, [1])
68 self.assertEqual(x, 42)
69 state.append(x)
70 self.assertEqual(state, [1, 42, 999])
71
72 def test_contextmanager_finally(self):
73 state = []
74 @contextmanager
75 def woohoo():
76 state.append(1)
77 try:
78 yield 42
79 finally:
80 state.append(999)
Florent Xicluna41fe6152010-04-02 18:52:12 +000081 with self.assertRaises(ZeroDivisionError):
Guido van Rossum1a5e21e2006-02-28 21:57:43 +000082 with woohoo() as x:
83 self.assertEqual(state, [1])
84 self.assertEqual(x, 42)
85 state.append(x)
86 raise ZeroDivisionError()
Guido van Rossum1a5e21e2006-02-28 21:57:43 +000087 self.assertEqual(state, [1, 42, 999])
88
Thomas Wouters49fd7fa2006-04-21 10:40:58 +000089 def test_contextmanager_no_reraise(self):
90 @contextmanager
91 def whee():
92 yield
Thomas Wouters477c8d52006-05-27 19:21:47 +000093 ctx = whee()
Thomas Wouters49fd7fa2006-04-21 10:40:58 +000094 ctx.__enter__()
95 # Calling __exit__ should not result in an exception
Benjamin Petersonc9c0f202009-06-30 23:06:06 +000096 self.assertFalse(ctx.__exit__(TypeError, TypeError("foo"), None))
Thomas Wouters49fd7fa2006-04-21 10:40:58 +000097
98 def test_contextmanager_trap_yield_after_throw(self):
99 @contextmanager
100 def whoo():
101 try:
102 yield
103 except:
104 yield
Thomas Wouters477c8d52006-05-27 19:21:47 +0000105 ctx = whoo()
Thomas Wouters49fd7fa2006-04-21 10:40:58 +0000106 ctx.__enter__()
107 self.assertRaises(
108 RuntimeError, ctx.__exit__, TypeError, TypeError("foo"), None
109 )
110
Guido van Rossum1a5e21e2006-02-28 21:57:43 +0000111 def test_contextmanager_except(self):
112 state = []
113 @contextmanager
114 def woohoo():
115 state.append(1)
116 try:
117 yield 42
Guido van Rossumb940e112007-01-10 16:19:56 +0000118 except ZeroDivisionError as e:
Guido van Rossum1a5e21e2006-02-28 21:57:43 +0000119 state.append(e.args[0])
120 self.assertEqual(state, [1, 42, 999])
121 with woohoo() as x:
122 self.assertEqual(state, [1])
123 self.assertEqual(x, 42)
124 state.append(x)
125 raise ZeroDivisionError(999)
126 self.assertEqual(state, [1, 42, 999])
127
Yury Selivanov8170e8c2015-05-09 11:44:30 -0400128 def test_contextmanager_except_stopiter(self):
129 stop_exc = StopIteration('spam')
130 @contextmanager
131 def woohoo():
132 yield
133 try:
Martin Panter7e3a91a2016-02-10 04:40:48 +0000134 with self.assertWarnsRegex(DeprecationWarning,
Yury Selivanov68333392015-05-22 11:16:47 -0400135 "StopIteration"):
136 with woohoo():
137 raise stop_exc
Yury Selivanov8170e8c2015-05-09 11:44:30 -0400138 except Exception as ex:
139 self.assertIs(ex, stop_exc)
140 else:
141 self.fail('StopIteration was suppressed')
142
143 def test_contextmanager_except_pep479(self):
144 code = """\
145from __future__ import generator_stop
146from contextlib import contextmanager
147@contextmanager
148def woohoo():
149 yield
150"""
151 locals = {}
152 exec(code, locals, locals)
153 woohoo = locals['woohoo']
154
155 stop_exc = StopIteration('spam')
156 try:
157 with woohoo():
158 raise stop_exc
159 except Exception as ex:
160 self.assertIs(ex, stop_exc)
161 else:
162 self.fail('StopIteration was suppressed')
163
svelankar00c75e92017-04-11 05:11:13 -0400164 def test_contextmanager_do_not_unchain_non_stopiteration_exceptions(self):
165 @contextmanager
166 def test_issue29692():
167 try:
168 yield
169 except Exception as exc:
170 raise RuntimeError('issue29692:Chained') from exc
171 try:
172 with test_issue29692():
173 raise ZeroDivisionError
174 except Exception as ex:
175 self.assertIs(type(ex), RuntimeError)
176 self.assertEqual(ex.args[0], 'issue29692:Chained')
177 self.assertIsInstance(ex.__cause__, ZeroDivisionError)
178
179 try:
180 with test_issue29692():
181 raise StopIteration('issue29692:Unchained')
182 except Exception as ex:
183 self.assertIs(type(ex), StopIteration)
184 self.assertEqual(ex.args[0], 'issue29692:Unchained')
185 self.assertIsNone(ex.__cause__)
186
R. David Murray378c0cf2010-02-24 01:46:21 +0000187 def _create_contextmanager_attribs(self):
Thomas Wouters49fd7fa2006-04-21 10:40:58 +0000188 def attribs(**kw):
189 def decorate(func):
190 for k,v in kw.items():
191 setattr(func,k,v)
192 return func
193 return decorate
194 @contextmanager
195 @attribs(foo='bar')
196 def baz(spam):
197 """Whee!"""
R. David Murray378c0cf2010-02-24 01:46:21 +0000198 return baz
199
200 def test_contextmanager_attribs(self):
201 baz = self._create_contextmanager_attribs()
Thomas Wouters49fd7fa2006-04-21 10:40:58 +0000202 self.assertEqual(baz.__name__,'baz')
203 self.assertEqual(baz.foo, 'bar')
R. David Murray378c0cf2010-02-24 01:46:21 +0000204
Nick Coghlan561eb5c2013-10-26 22:20:43 +1000205 @support.requires_docstrings
R. David Murray378c0cf2010-02-24 01:46:21 +0000206 def test_contextmanager_doc_attrib(self):
207 baz = self._create_contextmanager_attribs()
Thomas Wouters49fd7fa2006-04-21 10:40:58 +0000208 self.assertEqual(baz.__doc__, "Whee!")
209
Nick Coghlan561eb5c2013-10-26 22:20:43 +1000210 @support.requires_docstrings
211 def test_instance_docstring_given_cm_docstring(self):
212 baz = self._create_contextmanager_attribs()(None)
213 self.assertEqual(baz.__doc__, "Whee!")
214
Serhiy Storchaka101ff352015-06-28 17:06:07 +0300215 def test_keywords(self):
216 # Ensure no keyword arguments are inhibited
217 @contextmanager
218 def woohoo(self, func, args, kwds):
219 yield (self, func, args, kwds)
220 with woohoo(self=11, func=22, args=33, kwds=44) as target:
221 self.assertEqual(target, (11, 22, 33, 44))
222
Martin Teichmanndd0e0872018-01-28 05:17:46 +0100223 def test_nokeepref(self):
224 class A:
225 pass
226
227 @contextmanager
228 def woohoo(a, b):
229 a = weakref.ref(a)
230 b = weakref.ref(b)
231 self.assertIsNone(a())
232 self.assertIsNone(b())
233 yield
234
235 with woohoo(A(), b=A()):
236 pass
237
238 def test_param_errors(self):
239 @contextmanager
240 def woohoo(a, *, b):
241 yield
242
243 with self.assertRaises(TypeError):
244 woohoo()
245 with self.assertRaises(TypeError):
246 woohoo(3, 5)
247 with self.assertRaises(TypeError):
248 woohoo(b=3)
249
250 def test_recursive(self):
251 depth = 0
252 @contextmanager
253 def woohoo():
254 nonlocal depth
255 before = depth
256 depth += 1
257 yield
258 depth -= 1
259 self.assertEqual(depth, before)
260
261 @woohoo()
262 def recursive():
263 if depth < 10:
264 recursive()
265
266 recursive()
267 self.assertEqual(depth, 0)
268
Nick Coghlan561eb5c2013-10-26 22:20:43 +1000269
Guido van Rossum1a5e21e2006-02-28 21:57:43 +0000270class ClosingTestCase(unittest.TestCase):
271
Nick Coghlan561eb5c2013-10-26 22:20:43 +1000272 @support.requires_docstrings
Nick Coghlan059def52013-10-26 18:08:15 +1000273 def test_instance_docs(self):
274 # Issue 19330: ensure context manager instances have good docstrings
275 cm_docstring = closing.__doc__
276 obj = closing(None)
277 self.assertEqual(obj.__doc__, cm_docstring)
Guido van Rossum1a5e21e2006-02-28 21:57:43 +0000278
279 def test_closing(self):
280 state = []
281 class C:
282 def close(self):
283 state.append(1)
284 x = C()
285 self.assertEqual(state, [])
286 with closing(x) as y:
287 self.assertEqual(x, y)
288 self.assertEqual(state, [1])
289
290 def test_closing_error(self):
291 state = []
292 class C:
293 def close(self):
294 state.append(1)
295 x = C()
296 self.assertEqual(state, [])
Florent Xicluna41fe6152010-04-02 18:52:12 +0000297 with self.assertRaises(ZeroDivisionError):
Guido van Rossum1a5e21e2006-02-28 21:57:43 +0000298 with closing(x) as y:
299 self.assertEqual(x, y)
Florent Xicluna41fe6152010-04-02 18:52:12 +0000300 1 / 0
301 self.assertEqual(state, [1])
Guido van Rossum1a5e21e2006-02-28 21:57:43 +0000302
Jesse-Bakker0784a2e2017-11-23 01:23:28 +0100303
304class NullcontextTestCase(unittest.TestCase):
305 def test_nullcontext(self):
306 class C:
307 pass
308 c = C()
309 with nullcontext(c) as c_in:
310 self.assertIs(c_in, c)
311
312
Guido van Rossum1a5e21e2006-02-28 21:57:43 +0000313class FileContextTestCase(unittest.TestCase):
314
315 def testWithOpen(self):
316 tfn = tempfile.mktemp()
317 try:
318 f = None
319 with open(tfn, "w") as f:
Benjamin Petersonc9c0f202009-06-30 23:06:06 +0000320 self.assertFalse(f.closed)
Guido van Rossum1a5e21e2006-02-28 21:57:43 +0000321 f.write("Booh\n")
Benjamin Petersonc9c0f202009-06-30 23:06:06 +0000322 self.assertTrue(f.closed)
Guido van Rossum1a5e21e2006-02-28 21:57:43 +0000323 f = None
Florent Xicluna41fe6152010-04-02 18:52:12 +0000324 with self.assertRaises(ZeroDivisionError):
Guido van Rossum1a5e21e2006-02-28 21:57:43 +0000325 with open(tfn, "r") as f:
Benjamin Petersonc9c0f202009-06-30 23:06:06 +0000326 self.assertFalse(f.closed)
Guido van Rossum1a5e21e2006-02-28 21:57:43 +0000327 self.assertEqual(f.read(), "Booh\n")
Florent Xicluna41fe6152010-04-02 18:52:12 +0000328 1 / 0
329 self.assertTrue(f.closed)
Guido van Rossum1a5e21e2006-02-28 21:57:43 +0000330 finally:
Hai Shi96a6a6d2020-07-09 21:25:10 +0800331 os_helper.unlink(tfn)
Guido van Rossum1a5e21e2006-02-28 21:57:43 +0000332
333class LockContextTestCase(unittest.TestCase):
334
335 def boilerPlate(self, lock, locked):
Benjamin Petersonc9c0f202009-06-30 23:06:06 +0000336 self.assertFalse(locked())
Guido van Rossum1a5e21e2006-02-28 21:57:43 +0000337 with lock:
Benjamin Petersonc9c0f202009-06-30 23:06:06 +0000338 self.assertTrue(locked())
339 self.assertFalse(locked())
Florent Xicluna41fe6152010-04-02 18:52:12 +0000340 with self.assertRaises(ZeroDivisionError):
Guido van Rossum1a5e21e2006-02-28 21:57:43 +0000341 with lock:
Benjamin Petersonc9c0f202009-06-30 23:06:06 +0000342 self.assertTrue(locked())
Florent Xicluna41fe6152010-04-02 18:52:12 +0000343 1 / 0
344 self.assertFalse(locked())
Guido van Rossum1a5e21e2006-02-28 21:57:43 +0000345
346 def testWithLock(self):
347 lock = threading.Lock()
348 self.boilerPlate(lock, lock.locked)
349
350 def testWithRLock(self):
351 lock = threading.RLock()
352 self.boilerPlate(lock, lock._is_owned)
353
354 def testWithCondition(self):
355 lock = threading.Condition()
356 def locked():
357 return lock._is_owned()
358 self.boilerPlate(lock, locked)
359
360 def testWithSemaphore(self):
361 lock = threading.Semaphore()
362 def locked():
363 if lock.acquire(False):
364 lock.release()
365 return False
366 else:
367 return True
368 self.boilerPlate(lock, locked)
369
370 def testWithBoundedSemaphore(self):
371 lock = threading.BoundedSemaphore()
372 def locked():
373 if lock.acquire(False):
374 lock.release()
375 return False
376 else:
377 return True
378 self.boilerPlate(lock, locked)
379
Michael Foordb3a89842010-06-30 12:17:50 +0000380
381class mycontext(ContextDecorator):
Nick Coghlan059def52013-10-26 18:08:15 +1000382 """Example decoration-compatible context manager for testing"""
Michael Foordb3a89842010-06-30 12:17:50 +0000383 started = False
384 exc = None
385 catch = False
386
387 def __enter__(self):
388 self.started = True
389 return self
390
391 def __exit__(self, *exc):
392 self.exc = exc
393 return self.catch
394
395
396class TestContextDecorator(unittest.TestCase):
397
Nick Coghlan561eb5c2013-10-26 22:20:43 +1000398 @support.requires_docstrings
Nick Coghlan059def52013-10-26 18:08:15 +1000399 def test_instance_docs(self):
400 # Issue 19330: ensure context manager instances have good docstrings
401 cm_docstring = mycontext.__doc__
402 obj = mycontext()
403 self.assertEqual(obj.__doc__, cm_docstring)
404
Michael Foordb3a89842010-06-30 12:17:50 +0000405 def test_contextdecorator(self):
406 context = mycontext()
407 with context as result:
408 self.assertIs(result, context)
409 self.assertTrue(context.started)
410
411 self.assertEqual(context.exc, (None, None, None))
412
413
414 def test_contextdecorator_with_exception(self):
415 context = mycontext()
416
Ezio Melottied3a7d22010-12-01 02:32:32 +0000417 with self.assertRaisesRegex(NameError, 'foo'):
Michael Foordb3a89842010-06-30 12:17:50 +0000418 with context:
419 raise NameError('foo')
420 self.assertIsNotNone(context.exc)
421 self.assertIs(context.exc[0], NameError)
422
423 context = mycontext()
424 context.catch = True
425 with context:
426 raise NameError('foo')
427 self.assertIsNotNone(context.exc)
428 self.assertIs(context.exc[0], NameError)
429
430
431 def test_decorator(self):
432 context = mycontext()
433
434 @context
435 def test():
436 self.assertIsNone(context.exc)
437 self.assertTrue(context.started)
438 test()
439 self.assertEqual(context.exc, (None, None, None))
440
441
442 def test_decorator_with_exception(self):
443 context = mycontext()
444
445 @context
446 def test():
447 self.assertIsNone(context.exc)
448 self.assertTrue(context.started)
449 raise NameError('foo')
450
Ezio Melottied3a7d22010-12-01 02:32:32 +0000451 with self.assertRaisesRegex(NameError, 'foo'):
Michael Foordb3a89842010-06-30 12:17:50 +0000452 test()
453 self.assertIsNotNone(context.exc)
454 self.assertIs(context.exc[0], NameError)
455
456
457 def test_decorating_method(self):
458 context = mycontext()
459
460 class Test(object):
461
462 @context
463 def method(self, a, b, c=None):
464 self.a = a
465 self.b = b
466 self.c = c
467
468 # these tests are for argument passing when used as a decorator
469 test = Test()
470 test.method(1, 2)
471 self.assertEqual(test.a, 1)
472 self.assertEqual(test.b, 2)
473 self.assertEqual(test.c, None)
474
475 test = Test()
476 test.method('a', 'b', 'c')
477 self.assertEqual(test.a, 'a')
478 self.assertEqual(test.b, 'b')
479 self.assertEqual(test.c, 'c')
480
481 test = Test()
482 test.method(a=1, b=2)
483 self.assertEqual(test.a, 1)
484 self.assertEqual(test.b, 2)
485
486
487 def test_typo_enter(self):
488 class mycontext(ContextDecorator):
489 def __unter__(self):
490 pass
491 def __exit__(self, *exc):
492 pass
493
494 with self.assertRaises(AttributeError):
495 with mycontext():
496 pass
497
498
499 def test_typo_exit(self):
500 class mycontext(ContextDecorator):
501 def __enter__(self):
502 pass
503 def __uxit__(self, *exc):
504 pass
505
506 with self.assertRaises(AttributeError):
507 with mycontext():
508 pass
509
510
511 def test_contextdecorator_as_mixin(self):
512 class somecontext(object):
513 started = False
514 exc = None
515
516 def __enter__(self):
517 self.started = True
518 return self
519
520 def __exit__(self, *exc):
521 self.exc = exc
522
523 class mycontext(somecontext, ContextDecorator):
524 pass
525
526 context = mycontext()
527 @context
528 def test():
529 self.assertIsNone(context.exc)
530 self.assertTrue(context.started)
531 test()
532 self.assertEqual(context.exc, (None, None, None))
533
534
535 def test_contextmanager_as_decorator(self):
Michael Foordb3a89842010-06-30 12:17:50 +0000536 @contextmanager
537 def woohoo(y):
538 state.append(y)
539 yield
540 state.append(999)
541
Nick Coghlan0ded3e32011-05-05 23:49:25 +1000542 state = []
Michael Foordb3a89842010-06-30 12:17:50 +0000543 @woohoo(1)
544 def test(x):
545 self.assertEqual(state, [1])
546 state.append(x)
547 test('something')
548 self.assertEqual(state, [1, 'something', 999])
549
Nick Coghlan0ded3e32011-05-05 23:49:25 +1000550 # Issue #11647: Ensure the decorated function is 'reusable'
551 state = []
552 test('something else')
553 self.assertEqual(state, [1, 'something else', 999])
554
Michael Foordb3a89842010-06-30 12:17:50 +0000555
Ilya Kulakov1aa094f2018-01-25 12:51:18 -0800556class TestBaseExitStack:
557 exit_stack = None
Nick Coghlan3267a302012-05-21 22:54:43 +1000558
Nick Coghlan561eb5c2013-10-26 22:20:43 +1000559 @support.requires_docstrings
Nick Coghlan059def52013-10-26 18:08:15 +1000560 def test_instance_docs(self):
561 # Issue 19330: ensure context manager instances have good docstrings
Ilya Kulakov1aa094f2018-01-25 12:51:18 -0800562 cm_docstring = self.exit_stack.__doc__
563 obj = self.exit_stack()
Nick Coghlan059def52013-10-26 18:08:15 +1000564 self.assertEqual(obj.__doc__, cm_docstring)
565
Nick Coghlan3267a302012-05-21 22:54:43 +1000566 def test_no_resources(self):
Ilya Kulakov1aa094f2018-01-25 12:51:18 -0800567 with self.exit_stack():
Nick Coghlan3267a302012-05-21 22:54:43 +1000568 pass
569
570 def test_callback(self):
571 expected = [
572 ((), {}),
573 ((1,), {}),
574 ((1,2), {}),
575 ((), dict(example=1)),
576 ((1,), dict(example=1)),
577 ((1,2), dict(example=1)),
Serhiy Storchaka42a139e2019-04-01 09:16:35 +0300578 ((1,2), dict(self=3, callback=4)),
Nick Coghlan3267a302012-05-21 22:54:43 +1000579 ]
580 result = []
581 def _exit(*args, **kwds):
582 """Test metadata propagation"""
583 result.append((args, kwds))
Ilya Kulakov1aa094f2018-01-25 12:51:18 -0800584 with self.exit_stack() as stack:
Nick Coghlan3267a302012-05-21 22:54:43 +1000585 for args, kwds in reversed(expected):
586 if args and kwds:
587 f = stack.callback(_exit, *args, **kwds)
588 elif args:
589 f = stack.callback(_exit, *args)
590 elif kwds:
591 f = stack.callback(_exit, **kwds)
592 else:
593 f = stack.callback(_exit)
594 self.assertIs(f, _exit)
595 for wrapper in stack._exit_callbacks:
Ilya Kulakov1aa094f2018-01-25 12:51:18 -0800596 self.assertIs(wrapper[1].__wrapped__, _exit)
597 self.assertNotEqual(wrapper[1].__name__, _exit.__name__)
598 self.assertIsNone(wrapper[1].__doc__, _exit.__doc__)
Nick Coghlan3267a302012-05-21 22:54:43 +1000599 self.assertEqual(result, expected)
600
Serhiy Storchaka42a139e2019-04-01 09:16:35 +0300601 result = []
602 with self.exit_stack() as stack:
603 with self.assertRaises(TypeError):
604 stack.callback(arg=1)
605 with self.assertRaises(TypeError):
606 self.exit_stack.callback(arg=2)
Serhiy Storchaka142566c2019-06-05 18:22:31 +0300607 with self.assertRaises(TypeError):
Serhiy Storchaka42a139e2019-04-01 09:16:35 +0300608 stack.callback(callback=_exit, arg=3)
Serhiy Storchaka142566c2019-06-05 18:22:31 +0300609 self.assertEqual(result, [])
Serhiy Storchaka42a139e2019-04-01 09:16:35 +0300610
Nick Coghlan3267a302012-05-21 22:54:43 +1000611 def test_push(self):
612 exc_raised = ZeroDivisionError
613 def _expect_exc(exc_type, exc, exc_tb):
614 self.assertIs(exc_type, exc_raised)
615 def _suppress_exc(*exc_details):
616 return True
617 def _expect_ok(exc_type, exc, exc_tb):
618 self.assertIsNone(exc_type)
619 self.assertIsNone(exc)
620 self.assertIsNone(exc_tb)
621 class ExitCM(object):
622 def __init__(self, check_exc):
623 self.check_exc = check_exc
624 def __enter__(self):
625 self.fail("Should not be called!")
626 def __exit__(self, *exc_details):
627 self.check_exc(*exc_details)
Ilya Kulakov1aa094f2018-01-25 12:51:18 -0800628 with self.exit_stack() as stack:
Nick Coghlan3267a302012-05-21 22:54:43 +1000629 stack.push(_expect_ok)
Ilya Kulakov1aa094f2018-01-25 12:51:18 -0800630 self.assertIs(stack._exit_callbacks[-1][1], _expect_ok)
Nick Coghlan3267a302012-05-21 22:54:43 +1000631 cm = ExitCM(_expect_ok)
632 stack.push(cm)
Ilya Kulakov1aa094f2018-01-25 12:51:18 -0800633 self.assertIs(stack._exit_callbacks[-1][1].__self__, cm)
Nick Coghlan3267a302012-05-21 22:54:43 +1000634 stack.push(_suppress_exc)
Ilya Kulakov1aa094f2018-01-25 12:51:18 -0800635 self.assertIs(stack._exit_callbacks[-1][1], _suppress_exc)
Nick Coghlan3267a302012-05-21 22:54:43 +1000636 cm = ExitCM(_expect_exc)
637 stack.push(cm)
Ilya Kulakov1aa094f2018-01-25 12:51:18 -0800638 self.assertIs(stack._exit_callbacks[-1][1].__self__, cm)
Nick Coghlan3267a302012-05-21 22:54:43 +1000639 stack.push(_expect_exc)
Ilya Kulakov1aa094f2018-01-25 12:51:18 -0800640 self.assertIs(stack._exit_callbacks[-1][1], _expect_exc)
Nick Coghlan3267a302012-05-21 22:54:43 +1000641 stack.push(_expect_exc)
Ilya Kulakov1aa094f2018-01-25 12:51:18 -0800642 self.assertIs(stack._exit_callbacks[-1][1], _expect_exc)
Nick Coghlan3267a302012-05-21 22:54:43 +1000643 1/0
644
645 def test_enter_context(self):
646 class TestCM(object):
647 def __enter__(self):
648 result.append(1)
649 def __exit__(self, *exc_details):
650 result.append(3)
651
652 result = []
653 cm = TestCM()
Ilya Kulakov1aa094f2018-01-25 12:51:18 -0800654 with self.exit_stack() as stack:
Nick Coghlan3267a302012-05-21 22:54:43 +1000655 @stack.callback # Registered first => cleaned up last
656 def _exit():
657 result.append(4)
658 self.assertIsNotNone(_exit)
659 stack.enter_context(cm)
Ilya Kulakov1aa094f2018-01-25 12:51:18 -0800660 self.assertIs(stack._exit_callbacks[-1][1].__self__, cm)
Nick Coghlan3267a302012-05-21 22:54:43 +1000661 result.append(2)
662 self.assertEqual(result, [1, 2, 3, 4])
663
664 def test_close(self):
665 result = []
Ilya Kulakov1aa094f2018-01-25 12:51:18 -0800666 with self.exit_stack() as stack:
Nick Coghlan3267a302012-05-21 22:54:43 +1000667 @stack.callback
668 def _exit():
669 result.append(1)
670 self.assertIsNotNone(_exit)
671 stack.close()
672 result.append(2)
673 self.assertEqual(result, [1, 2])
674
675 def test_pop_all(self):
676 result = []
Ilya Kulakov1aa094f2018-01-25 12:51:18 -0800677 with self.exit_stack() as stack:
Nick Coghlan3267a302012-05-21 22:54:43 +1000678 @stack.callback
679 def _exit():
680 result.append(3)
681 self.assertIsNotNone(_exit)
682 new_stack = stack.pop_all()
683 result.append(1)
684 result.append(2)
685 new_stack.close()
686 self.assertEqual(result, [1, 2, 3])
687
Nick Coghlanc73e8c22012-05-31 23:49:26 +1000688 def test_exit_raise(self):
689 with self.assertRaises(ZeroDivisionError):
Ilya Kulakov1aa094f2018-01-25 12:51:18 -0800690 with self.exit_stack() as stack:
Nick Coghlanc73e8c22012-05-31 23:49:26 +1000691 stack.push(lambda *exc: False)
692 1/0
693
694 def test_exit_suppress(self):
Ilya Kulakov1aa094f2018-01-25 12:51:18 -0800695 with self.exit_stack() as stack:
Nick Coghlanc73e8c22012-05-31 23:49:26 +1000696 stack.push(lambda *exc: True)
697 1/0
698
699 def test_exit_exception_chaining_reference(self):
700 # Sanity check to make sure that ExitStack chaining matches
701 # actual nested with statements
702 class RaiseExc:
703 def __init__(self, exc):
704 self.exc = exc
705 def __enter__(self):
706 return self
707 def __exit__(self, *exc_details):
708 raise self.exc
709
Nick Coghlan77452fc2012-06-01 22:48:32 +1000710 class RaiseExcWithContext:
711 def __init__(self, outer, inner):
712 self.outer = outer
713 self.inner = inner
714 def __enter__(self):
715 return self
716 def __exit__(self, *exc_details):
717 try:
718 raise self.inner
719 except:
720 raise self.outer
721
Nick Coghlanc73e8c22012-05-31 23:49:26 +1000722 class SuppressExc:
723 def __enter__(self):
724 return self
725 def __exit__(self, *exc_details):
726 type(self).saved_details = exc_details
727 return True
728
729 try:
730 with RaiseExc(IndexError):
Nick Coghlan77452fc2012-06-01 22:48:32 +1000731 with RaiseExcWithContext(KeyError, AttributeError):
732 with SuppressExc():
733 with RaiseExc(ValueError):
734 1 / 0
Nick Coghlanc73e8c22012-05-31 23:49:26 +1000735 except IndexError as exc:
736 self.assertIsInstance(exc.__context__, KeyError)
737 self.assertIsInstance(exc.__context__.__context__, AttributeError)
738 # Inner exceptions were suppressed
739 self.assertIsNone(exc.__context__.__context__.__context__)
740 else:
741 self.fail("Expected IndexError, but no exception was raised")
742 # Check the inner exceptions
743 inner_exc = SuppressExc.saved_details[1]
744 self.assertIsInstance(inner_exc, ValueError)
745 self.assertIsInstance(inner_exc.__context__, ZeroDivisionError)
746
747 def test_exit_exception_chaining(self):
748 # Ensure exception chaining matches the reference behaviour
749 def raise_exc(exc):
750 raise exc
751
752 saved_details = None
753 def suppress_exc(*exc_details):
754 nonlocal saved_details
755 saved_details = exc_details
756 return True
757
758 try:
Ilya Kulakov1aa094f2018-01-25 12:51:18 -0800759 with self.exit_stack() as stack:
Nick Coghlanc73e8c22012-05-31 23:49:26 +1000760 stack.callback(raise_exc, IndexError)
761 stack.callback(raise_exc, KeyError)
762 stack.callback(raise_exc, AttributeError)
763 stack.push(suppress_exc)
764 stack.callback(raise_exc, ValueError)
765 1 / 0
766 except IndexError as exc:
767 self.assertIsInstance(exc.__context__, KeyError)
768 self.assertIsInstance(exc.__context__.__context__, AttributeError)
Nick Coghlan77452fc2012-06-01 22:48:32 +1000769 # Inner exceptions were suppressed
770 self.assertIsNone(exc.__context__.__context__.__context__)
Nick Coghlanc73e8c22012-05-31 23:49:26 +1000771 else:
772 self.fail("Expected IndexError, but no exception was raised")
773 # Check the inner exceptions
774 inner_exc = saved_details[1]
775 self.assertIsInstance(inner_exc, ValueError)
776 self.assertIsInstance(inner_exc.__context__, ZeroDivisionError)
777
Nick Coghlan1a33b2f2013-10-01 23:24:56 +1000778 def test_exit_exception_non_suppressing(self):
779 # http://bugs.python.org/issue19092
780 def raise_exc(exc):
781 raise exc
782
783 def suppress_exc(*exc_details):
784 return True
785
786 try:
Ilya Kulakov1aa094f2018-01-25 12:51:18 -0800787 with self.exit_stack() as stack:
Nick Coghlan1a33b2f2013-10-01 23:24:56 +1000788 stack.callback(lambda: None)
789 stack.callback(raise_exc, IndexError)
790 except Exception as exc:
791 self.assertIsInstance(exc, IndexError)
792 else:
793 self.fail("Expected IndexError, but no exception was raised")
794
795 try:
Ilya Kulakov1aa094f2018-01-25 12:51:18 -0800796 with self.exit_stack() as stack:
Nick Coghlan1a33b2f2013-10-01 23:24:56 +1000797 stack.callback(raise_exc, KeyError)
798 stack.push(suppress_exc)
799 stack.callback(raise_exc, IndexError)
800 except Exception as exc:
801 self.assertIsInstance(exc, KeyError)
802 else:
803 self.fail("Expected KeyError, but no exception was raised")
804
Nick Coghlan09761e72014-01-22 22:24:46 +1000805 def test_exit_exception_with_correct_context(self):
806 # http://bugs.python.org/issue20317
807 @contextmanager
Nick Coghlanadd94c92014-01-24 23:05:45 +1000808 def gets_the_context_right(exc):
Nick Coghlan09761e72014-01-22 22:24:46 +1000809 try:
Nick Coghlanadd94c92014-01-24 23:05:45 +1000810 yield
Nick Coghlan09761e72014-01-22 22:24:46 +1000811 finally:
Nick Coghlanadd94c92014-01-24 23:05:45 +1000812 raise exc
813
814 exc1 = Exception(1)
815 exc2 = Exception(2)
816 exc3 = Exception(3)
817 exc4 = Exception(4)
Nick Coghlan09761e72014-01-22 22:24:46 +1000818
819 # The contextmanager already fixes the context, so prior to the
820 # fix, ExitStack would try to fix it *again* and get into an
821 # infinite self-referential loop
822 try:
Ilya Kulakov1aa094f2018-01-25 12:51:18 -0800823 with self.exit_stack() as stack:
Nick Coghlanadd94c92014-01-24 23:05:45 +1000824 stack.enter_context(gets_the_context_right(exc4))
825 stack.enter_context(gets_the_context_right(exc3))
826 stack.enter_context(gets_the_context_right(exc2))
827 raise exc1
828 except Exception as exc:
829 self.assertIs(exc, exc4)
830 self.assertIs(exc.__context__, exc3)
831 self.assertIs(exc.__context__.__context__, exc2)
832 self.assertIs(exc.__context__.__context__.__context__, exc1)
833 self.assertIsNone(
834 exc.__context__.__context__.__context__.__context__)
835
836 def test_exit_exception_with_existing_context(self):
837 # Addresses a lack of test coverage discovered after checking in a
838 # fix for issue 20317 that still contained debugging code.
839 def raise_nested(inner_exc, outer_exc):
840 try:
841 raise inner_exc
842 finally:
843 raise outer_exc
844 exc1 = Exception(1)
845 exc2 = Exception(2)
846 exc3 = Exception(3)
847 exc4 = Exception(4)
848 exc5 = Exception(5)
849 try:
Ilya Kulakov1aa094f2018-01-25 12:51:18 -0800850 with self.exit_stack() as stack:
Nick Coghlanadd94c92014-01-24 23:05:45 +1000851 stack.callback(raise_nested, exc4, exc5)
852 stack.callback(raise_nested, exc2, exc3)
853 raise exc1
854 except Exception as exc:
855 self.assertIs(exc, exc5)
856 self.assertIs(exc.__context__, exc4)
857 self.assertIs(exc.__context__.__context__, exc3)
858 self.assertIs(exc.__context__.__context__.__context__, exc2)
859 self.assertIs(
860 exc.__context__.__context__.__context__.__context__, exc1)
861 self.assertIsNone(
862 exc.__context__.__context__.__context__.__context__.__context__)
863
Nick Coghlan1a33b2f2013-10-01 23:24:56 +1000864 def test_body_exception_suppress(self):
865 def suppress_exc(*exc_details):
866 return True
867 try:
Ilya Kulakov1aa094f2018-01-25 12:51:18 -0800868 with self.exit_stack() as stack:
Nick Coghlan1a33b2f2013-10-01 23:24:56 +1000869 stack.push(suppress_exc)
870 1/0
871 except IndexError as exc:
872 self.fail("Expected no exception, got IndexError")
873
Nick Coghlanc73e8c22012-05-31 23:49:26 +1000874 def test_exit_exception_chaining_suppress(self):
Ilya Kulakov1aa094f2018-01-25 12:51:18 -0800875 with self.exit_stack() as stack:
Nick Coghlanc73e8c22012-05-31 23:49:26 +1000876 stack.push(lambda *exc: True)
877 stack.push(lambda *exc: 1/0)
878 stack.push(lambda *exc: {}[1])
879
Nick Coghlana5bd2a12012-06-01 00:00:38 +1000880 def test_excessive_nesting(self):
881 # The original implementation would die with RecursionError here
Ilya Kulakov1aa094f2018-01-25 12:51:18 -0800882 with self.exit_stack() as stack:
Nick Coghlana5bd2a12012-06-01 00:00:38 +1000883 for i in range(10000):
884 stack.callback(int)
885
Nick Coghlan3267a302012-05-21 22:54:43 +1000886 def test_instance_bypass(self):
887 class Example(object): pass
888 cm = Example()
889 cm.__exit__ = object()
Ilya Kulakov1aa094f2018-01-25 12:51:18 -0800890 stack = self.exit_stack()
Nick Coghlan3267a302012-05-21 22:54:43 +1000891 self.assertRaises(AttributeError, stack.enter_context, cm)
892 stack.push(cm)
Ilya Kulakov1aa094f2018-01-25 12:51:18 -0800893 self.assertIs(stack._exit_callbacks[-1][1], cm)
Nick Coghlan3267a302012-05-21 22:54:43 +1000894
Gregory P. Smithba2ecd62016-06-14 09:19:20 -0700895 def test_dont_reraise_RuntimeError(self):
Serhiy Storchakace1a9f32016-06-20 05:29:54 +0300896 # https://bugs.python.org/issue27122
Gregory P. Smithba2ecd62016-06-14 09:19:20 -0700897 class UniqueException(Exception): pass
Serhiy Storchakace1a9f32016-06-20 05:29:54 +0300898 class UniqueRuntimeError(RuntimeError): pass
Gregory P. Smithba2ecd62016-06-14 09:19:20 -0700899
900 @contextmanager
901 def second():
902 try:
903 yield 1
904 except Exception as exc:
905 raise UniqueException("new exception") from exc
906
907 @contextmanager
908 def first():
909 try:
910 yield 1
911 except Exception as exc:
912 raise exc
913
Serhiy Storchakace1a9f32016-06-20 05:29:54 +0300914 # The UniqueRuntimeError should be caught by second()'s exception
Gregory P. Smithba2ecd62016-06-14 09:19:20 -0700915 # handler which chain raised a new UniqueException.
916 with self.assertRaises(UniqueException) as err_ctx:
Ilya Kulakov1aa094f2018-01-25 12:51:18 -0800917 with self.exit_stack() as es_ctx:
Gregory P. Smithba2ecd62016-06-14 09:19:20 -0700918 es_ctx.enter_context(second())
919 es_ctx.enter_context(first())
Serhiy Storchakace1a9f32016-06-20 05:29:54 +0300920 raise UniqueRuntimeError("please no infinite loop.")
Gregory P. Smithba2ecd62016-06-14 09:19:20 -0700921
Serhiy Storchakace1a9f32016-06-20 05:29:54 +0300922 exc = err_ctx.exception
923 self.assertIsInstance(exc, UniqueException)
924 self.assertIsInstance(exc.__context__, UniqueRuntimeError)
925 self.assertIsNone(exc.__context__.__context__)
926 self.assertIsNone(exc.__context__.__cause__)
927 self.assertIs(exc.__cause__, exc.__context__)
Gregory P. Smithba2ecd62016-06-14 09:19:20 -0700928
Berker Peksagbb44fe02014-11-28 23:28:06 +0200929
Ilya Kulakov1aa094f2018-01-25 12:51:18 -0800930class TestExitStack(TestBaseExitStack, unittest.TestCase):
931 exit_stack = ExitStack
932
933
Berker Peksagbb44fe02014-11-28 23:28:06 +0200934class TestRedirectStream:
935
936 redirect_stream = None
937 orig_stream = None
Raymond Hettinger088cbf22013-10-10 00:46:57 -0700938
Nick Coghlan561eb5c2013-10-26 22:20:43 +1000939 @support.requires_docstrings
Nick Coghlan059def52013-10-26 18:08:15 +1000940 def test_instance_docs(self):
941 # Issue 19330: ensure context manager instances have good docstrings
Berker Peksagbb44fe02014-11-28 23:28:06 +0200942 cm_docstring = self.redirect_stream.__doc__
943 obj = self.redirect_stream(None)
Nick Coghlan059def52013-10-26 18:08:15 +1000944 self.assertEqual(obj.__doc__, cm_docstring)
945
Nick Coghlan8e113b42013-11-03 17:00:51 +1000946 def test_no_redirect_in_init(self):
Berker Peksagbb44fe02014-11-28 23:28:06 +0200947 orig_stdout = getattr(sys, self.orig_stream)
948 self.redirect_stream(None)
949 self.assertIs(getattr(sys, self.orig_stream), orig_stdout)
Nick Coghlan8e113b42013-11-03 17:00:51 +1000950
Raymond Hettinger088cbf22013-10-10 00:46:57 -0700951 def test_redirect_to_string_io(self):
952 f = io.StringIO()
Nick Coghlan0ddaed32013-10-26 16:37:47 +1000953 msg = "Consider an API like help(), which prints directly to stdout"
Berker Peksagbb44fe02014-11-28 23:28:06 +0200954 orig_stdout = getattr(sys, self.orig_stream)
955 with self.redirect_stream(f):
956 print(msg, file=getattr(sys, self.orig_stream))
957 self.assertIs(getattr(sys, self.orig_stream), orig_stdout)
Nick Coghlan0ddaed32013-10-26 16:37:47 +1000958 s = f.getvalue().strip()
959 self.assertEqual(s, msg)
Nick Coghlan3267a302012-05-21 22:54:43 +1000960
Nick Coghlan8608d262013-10-20 00:30:51 +1000961 def test_enter_result_is_target(self):
962 f = io.StringIO()
Berker Peksagbb44fe02014-11-28 23:28:06 +0200963 with self.redirect_stream(f) as enter_result:
Nick Coghlan8608d262013-10-20 00:30:51 +1000964 self.assertIs(enter_result, f)
965
966 def test_cm_is_reusable(self):
967 f = io.StringIO()
Berker Peksagbb44fe02014-11-28 23:28:06 +0200968 write_to_f = self.redirect_stream(f)
969 orig_stdout = getattr(sys, self.orig_stream)
Nick Coghlan8608d262013-10-20 00:30:51 +1000970 with write_to_f:
Berker Peksagbb44fe02014-11-28 23:28:06 +0200971 print("Hello", end=" ", file=getattr(sys, self.orig_stream))
Nick Coghlan8608d262013-10-20 00:30:51 +1000972 with write_to_f:
Berker Peksagbb44fe02014-11-28 23:28:06 +0200973 print("World!", file=getattr(sys, self.orig_stream))
974 self.assertIs(getattr(sys, self.orig_stream), orig_stdout)
Nick Coghlan8608d262013-10-20 00:30:51 +1000975 s = f.getvalue()
976 self.assertEqual(s, "Hello World!\n")
977
Nick Coghlan8e113b42013-11-03 17:00:51 +1000978 def test_cm_is_reentrant(self):
Nick Coghlan8608d262013-10-20 00:30:51 +1000979 f = io.StringIO()
Berker Peksagbb44fe02014-11-28 23:28:06 +0200980 write_to_f = self.redirect_stream(f)
981 orig_stdout = getattr(sys, self.orig_stream)
Nick Coghlan8e113b42013-11-03 17:00:51 +1000982 with write_to_f:
Berker Peksagbb44fe02014-11-28 23:28:06 +0200983 print("Hello", end=" ", file=getattr(sys, self.orig_stream))
Nick Coghlan8608d262013-10-20 00:30:51 +1000984 with write_to_f:
Berker Peksagbb44fe02014-11-28 23:28:06 +0200985 print("World!", file=getattr(sys, self.orig_stream))
986 self.assertIs(getattr(sys, self.orig_stream), orig_stdout)
Nick Coghlan8e113b42013-11-03 17:00:51 +1000987 s = f.getvalue()
988 self.assertEqual(s, "Hello World!\n")
Nick Coghlan8608d262013-10-20 00:30:51 +1000989
990
Berker Peksagbb44fe02014-11-28 23:28:06 +0200991class TestRedirectStdout(TestRedirectStream, unittest.TestCase):
992
993 redirect_stream = redirect_stdout
994 orig_stream = "stdout"
995
996
997class TestRedirectStderr(TestRedirectStream, unittest.TestCase):
998
999 redirect_stream = redirect_stderr
1000 orig_stream = "stderr"
1001
1002
Nick Coghlan240f86d2013-10-17 23:40:57 +10001003class TestSuppress(unittest.TestCase):
1004
Nick Coghlan561eb5c2013-10-26 22:20:43 +10001005 @support.requires_docstrings
Nick Coghlan059def52013-10-26 18:08:15 +10001006 def test_instance_docs(self):
1007 # Issue 19330: ensure context manager instances have good docstrings
1008 cm_docstring = suppress.__doc__
1009 obj = suppress()
1010 self.assertEqual(obj.__doc__, cm_docstring)
1011
Nick Coghlan8608d262013-10-20 00:30:51 +10001012 def test_no_result_from_enter(self):
1013 with suppress(ValueError) as enter_result:
1014 self.assertIsNone(enter_result)
Nick Coghlan240f86d2013-10-17 23:40:57 +10001015
Nick Coghlan8608d262013-10-20 00:30:51 +10001016 def test_no_exception(self):
Nick Coghlan240f86d2013-10-17 23:40:57 +10001017 with suppress(ValueError):
1018 self.assertEqual(pow(2, 5), 32)
1019
1020 def test_exact_exception(self):
Nick Coghlan240f86d2013-10-17 23:40:57 +10001021 with suppress(TypeError):
1022 len(5)
1023
Nick Coghlan059def52013-10-26 18:08:15 +10001024 def test_exception_hierarchy(self):
1025 with suppress(LookupError):
1026 'Hello'[50]
1027
1028 def test_other_exception(self):
1029 with self.assertRaises(ZeroDivisionError):
1030 with suppress(TypeError):
1031 1/0
1032
1033 def test_no_args(self):
1034 with self.assertRaises(ZeroDivisionError):
1035 with suppress():
1036 1/0
1037
Nick Coghlan240f86d2013-10-17 23:40:57 +10001038 def test_multiple_exception_args(self):
Nick Coghlan8608d262013-10-20 00:30:51 +10001039 with suppress(ZeroDivisionError, TypeError):
1040 1/0
Nick Coghlan240f86d2013-10-17 23:40:57 +10001041 with suppress(ZeroDivisionError, TypeError):
1042 len(5)
1043
Nick Coghlan8608d262013-10-20 00:30:51 +10001044 def test_cm_is_reentrant(self):
1045 ignore_exceptions = suppress(Exception)
1046 with ignore_exceptions:
1047 pass
1048 with ignore_exceptions:
1049 len(5)
1050 with ignore_exceptions:
Nick Coghlan8608d262013-10-20 00:30:51 +10001051 with ignore_exceptions: # Check nested usage
1052 len(5)
Martin Panter7c6420a2015-10-10 11:04:44 +00001053 outer_continued = True
1054 1/0
1055 self.assertTrue(outer_continued)
Nick Coghlan8608d262013-10-20 00:30:51 +10001056
Guido van Rossum1a5e21e2006-02-28 21:57:43 +00001057if __name__ == "__main__":
Brett Cannon3e9a9ae2013-06-12 21:25:59 -04001058 unittest.main()