blob: 1a5e6edad9b28f22e14a22225b7d4f06a026b679 [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
Guido van Rossum1a5e21e2006-02-28 21:57:43 +000010
Florent Xicluna41fe6152010-04-02 18:52:12 +000011
Brett Cannon9e080e02016-04-08 12:15:27 -070012class TestAbstractContextManager(unittest.TestCase):
13
14 def test_enter(self):
15 class DefaultEnter(AbstractContextManager):
16 def __exit__(self, *args):
17 super().__exit__(*args)
18
19 manager = DefaultEnter()
20 self.assertIs(manager.__enter__(), manager)
21
22 def test_exit_is_abstract(self):
23 class MissingExit(AbstractContextManager):
24 pass
25
26 with self.assertRaises(TypeError):
27 MissingExit()
28
29 def test_structural_subclassing(self):
30 class ManagerFromScratch:
31 def __enter__(self):
32 return self
33 def __exit__(self, exc_type, exc_value, traceback):
34 return None
35
36 self.assertTrue(issubclass(ManagerFromScratch, AbstractContextManager))
37
38 class DefaultEnter(AbstractContextManager):
39 def __exit__(self, *args):
40 super().__exit__(*args)
41
42 self.assertTrue(issubclass(DefaultEnter, AbstractContextManager))
43
Jelle Zijlstra57161aa2017-06-09 08:21:47 -070044 class NoEnter(ManagerFromScratch):
45 __enter__ = None
46
47 self.assertFalse(issubclass(NoEnter, AbstractContextManager))
48
49 class NoExit(ManagerFromScratch):
50 __exit__ = None
51
52 self.assertFalse(issubclass(NoExit, AbstractContextManager))
53
Brett Cannon9e080e02016-04-08 12:15:27 -070054
Guido van Rossum1a5e21e2006-02-28 21:57:43 +000055class ContextManagerTestCase(unittest.TestCase):
56
57 def test_contextmanager_plain(self):
58 state = []
59 @contextmanager
60 def woohoo():
61 state.append(1)
62 yield 42
63 state.append(999)
64 with woohoo() as x:
65 self.assertEqual(state, [1])
66 self.assertEqual(x, 42)
67 state.append(x)
68 self.assertEqual(state, [1, 42, 999])
69
70 def test_contextmanager_finally(self):
71 state = []
72 @contextmanager
73 def woohoo():
74 state.append(1)
75 try:
76 yield 42
77 finally:
78 state.append(999)
Florent Xicluna41fe6152010-04-02 18:52:12 +000079 with self.assertRaises(ZeroDivisionError):
Guido van Rossum1a5e21e2006-02-28 21:57:43 +000080 with woohoo() as x:
81 self.assertEqual(state, [1])
82 self.assertEqual(x, 42)
83 state.append(x)
84 raise ZeroDivisionError()
Guido van Rossum1a5e21e2006-02-28 21:57:43 +000085 self.assertEqual(state, [1, 42, 999])
86
Thomas Wouters49fd7fa2006-04-21 10:40:58 +000087 def test_contextmanager_no_reraise(self):
88 @contextmanager
89 def whee():
90 yield
Thomas Wouters477c8d52006-05-27 19:21:47 +000091 ctx = whee()
Thomas Wouters49fd7fa2006-04-21 10:40:58 +000092 ctx.__enter__()
93 # Calling __exit__ should not result in an exception
Benjamin Petersonc9c0f202009-06-30 23:06:06 +000094 self.assertFalse(ctx.__exit__(TypeError, TypeError("foo"), None))
Thomas Wouters49fd7fa2006-04-21 10:40:58 +000095
96 def test_contextmanager_trap_yield_after_throw(self):
97 @contextmanager
98 def whoo():
99 try:
100 yield
101 except:
102 yield
Thomas Wouters477c8d52006-05-27 19:21:47 +0000103 ctx = whoo()
Thomas Wouters49fd7fa2006-04-21 10:40:58 +0000104 ctx.__enter__()
105 self.assertRaises(
106 RuntimeError, ctx.__exit__, TypeError, TypeError("foo"), None
107 )
108
Guido van Rossum1a5e21e2006-02-28 21:57:43 +0000109 def test_contextmanager_except(self):
110 state = []
111 @contextmanager
112 def woohoo():
113 state.append(1)
114 try:
115 yield 42
Guido van Rossumb940e112007-01-10 16:19:56 +0000116 except ZeroDivisionError as e:
Guido van Rossum1a5e21e2006-02-28 21:57:43 +0000117 state.append(e.args[0])
118 self.assertEqual(state, [1, 42, 999])
119 with woohoo() as x:
120 self.assertEqual(state, [1])
121 self.assertEqual(x, 42)
122 state.append(x)
123 raise ZeroDivisionError(999)
124 self.assertEqual(state, [1, 42, 999])
125
Yury Selivanov8170e8c2015-05-09 11:44:30 -0400126 def test_contextmanager_except_stopiter(self):
127 stop_exc = StopIteration('spam')
128 @contextmanager
129 def woohoo():
130 yield
131 try:
Martin Panter7e3a91a2016-02-10 04:40:48 +0000132 with self.assertWarnsRegex(DeprecationWarning,
Yury Selivanov68333392015-05-22 11:16:47 -0400133 "StopIteration"):
134 with woohoo():
135 raise stop_exc
Yury Selivanov8170e8c2015-05-09 11:44:30 -0400136 except Exception as ex:
137 self.assertIs(ex, stop_exc)
138 else:
139 self.fail('StopIteration was suppressed')
140
141 def test_contextmanager_except_pep479(self):
142 code = """\
143from __future__ import generator_stop
144from contextlib import contextmanager
145@contextmanager
146def woohoo():
147 yield
148"""
149 locals = {}
150 exec(code, locals, locals)
151 woohoo = locals['woohoo']
152
153 stop_exc = StopIteration('spam')
154 try:
155 with woohoo():
156 raise stop_exc
157 except Exception as ex:
158 self.assertIs(ex, stop_exc)
159 else:
160 self.fail('StopIteration was suppressed')
161
svelankar00c75e92017-04-11 05:11:13 -0400162 def test_contextmanager_do_not_unchain_non_stopiteration_exceptions(self):
163 @contextmanager
164 def test_issue29692():
165 try:
166 yield
167 except Exception as exc:
168 raise RuntimeError('issue29692:Chained') from exc
169 try:
170 with test_issue29692():
171 raise ZeroDivisionError
172 except Exception as ex:
173 self.assertIs(type(ex), RuntimeError)
174 self.assertEqual(ex.args[0], 'issue29692:Chained')
175 self.assertIsInstance(ex.__cause__, ZeroDivisionError)
176
177 try:
178 with test_issue29692():
179 raise StopIteration('issue29692:Unchained')
180 except Exception as ex:
181 self.assertIs(type(ex), StopIteration)
182 self.assertEqual(ex.args[0], 'issue29692:Unchained')
183 self.assertIsNone(ex.__cause__)
184
R. David Murray378c0cf2010-02-24 01:46:21 +0000185 def _create_contextmanager_attribs(self):
Thomas Wouters49fd7fa2006-04-21 10:40:58 +0000186 def attribs(**kw):
187 def decorate(func):
188 for k,v in kw.items():
189 setattr(func,k,v)
190 return func
191 return decorate
192 @contextmanager
193 @attribs(foo='bar')
194 def baz(spam):
195 """Whee!"""
R. David Murray378c0cf2010-02-24 01:46:21 +0000196 return baz
197
198 def test_contextmanager_attribs(self):
199 baz = self._create_contextmanager_attribs()
Thomas Wouters49fd7fa2006-04-21 10:40:58 +0000200 self.assertEqual(baz.__name__,'baz')
201 self.assertEqual(baz.foo, 'bar')
R. David Murray378c0cf2010-02-24 01:46:21 +0000202
Nick Coghlan561eb5c2013-10-26 22:20:43 +1000203 @support.requires_docstrings
R. David Murray378c0cf2010-02-24 01:46:21 +0000204 def test_contextmanager_doc_attrib(self):
205 baz = self._create_contextmanager_attribs()
Thomas Wouters49fd7fa2006-04-21 10:40:58 +0000206 self.assertEqual(baz.__doc__, "Whee!")
207
Nick Coghlan561eb5c2013-10-26 22:20:43 +1000208 @support.requires_docstrings
209 def test_instance_docstring_given_cm_docstring(self):
210 baz = self._create_contextmanager_attribs()(None)
211 self.assertEqual(baz.__doc__, "Whee!")
212
Serhiy Storchaka101ff352015-06-28 17:06:07 +0300213 def test_keywords(self):
214 # Ensure no keyword arguments are inhibited
215 @contextmanager
216 def woohoo(self, func, args, kwds):
217 yield (self, func, args, kwds)
218 with woohoo(self=11, func=22, args=33, kwds=44) as target:
219 self.assertEqual(target, (11, 22, 33, 44))
220
Nick Coghlan561eb5c2013-10-26 22:20:43 +1000221
Guido van Rossum1a5e21e2006-02-28 21:57:43 +0000222class ClosingTestCase(unittest.TestCase):
223
Nick Coghlan561eb5c2013-10-26 22:20:43 +1000224 @support.requires_docstrings
Nick Coghlan059def52013-10-26 18:08:15 +1000225 def test_instance_docs(self):
226 # Issue 19330: ensure context manager instances have good docstrings
227 cm_docstring = closing.__doc__
228 obj = closing(None)
229 self.assertEqual(obj.__doc__, cm_docstring)
Guido van Rossum1a5e21e2006-02-28 21:57:43 +0000230
231 def test_closing(self):
232 state = []
233 class C:
234 def close(self):
235 state.append(1)
236 x = C()
237 self.assertEqual(state, [])
238 with closing(x) as y:
239 self.assertEqual(x, y)
240 self.assertEqual(state, [1])
241
242 def test_closing_error(self):
243 state = []
244 class C:
245 def close(self):
246 state.append(1)
247 x = C()
248 self.assertEqual(state, [])
Florent Xicluna41fe6152010-04-02 18:52:12 +0000249 with self.assertRaises(ZeroDivisionError):
Guido van Rossum1a5e21e2006-02-28 21:57:43 +0000250 with closing(x) as y:
251 self.assertEqual(x, y)
Florent Xicluna41fe6152010-04-02 18:52:12 +0000252 1 / 0
253 self.assertEqual(state, [1])
Guido van Rossum1a5e21e2006-02-28 21:57:43 +0000254
Jesse-Bakker0784a2e2017-11-23 01:23:28 +0100255
256class NullcontextTestCase(unittest.TestCase):
257 def test_nullcontext(self):
258 class C:
259 pass
260 c = C()
261 with nullcontext(c) as c_in:
262 self.assertIs(c_in, c)
263
264
Guido van Rossum1a5e21e2006-02-28 21:57:43 +0000265class FileContextTestCase(unittest.TestCase):
266
267 def testWithOpen(self):
268 tfn = tempfile.mktemp()
269 try:
270 f = None
271 with open(tfn, "w") as f:
Benjamin Petersonc9c0f202009-06-30 23:06:06 +0000272 self.assertFalse(f.closed)
Guido van Rossum1a5e21e2006-02-28 21:57:43 +0000273 f.write("Booh\n")
Benjamin Petersonc9c0f202009-06-30 23:06:06 +0000274 self.assertTrue(f.closed)
Guido van Rossum1a5e21e2006-02-28 21:57:43 +0000275 f = None
Florent Xicluna41fe6152010-04-02 18:52:12 +0000276 with self.assertRaises(ZeroDivisionError):
Guido van Rossum1a5e21e2006-02-28 21:57:43 +0000277 with open(tfn, "r") as f:
Benjamin Petersonc9c0f202009-06-30 23:06:06 +0000278 self.assertFalse(f.closed)
Guido van Rossum1a5e21e2006-02-28 21:57:43 +0000279 self.assertEqual(f.read(), "Booh\n")
Florent Xicluna41fe6152010-04-02 18:52:12 +0000280 1 / 0
281 self.assertTrue(f.closed)
Guido van Rossum1a5e21e2006-02-28 21:57:43 +0000282 finally:
Florent Xicluna41fe6152010-04-02 18:52:12 +0000283 support.unlink(tfn)
Guido van Rossum1a5e21e2006-02-28 21:57:43 +0000284
285class LockContextTestCase(unittest.TestCase):
286
287 def boilerPlate(self, lock, locked):
Benjamin Petersonc9c0f202009-06-30 23:06:06 +0000288 self.assertFalse(locked())
Guido van Rossum1a5e21e2006-02-28 21:57:43 +0000289 with lock:
Benjamin Petersonc9c0f202009-06-30 23:06:06 +0000290 self.assertTrue(locked())
291 self.assertFalse(locked())
Florent Xicluna41fe6152010-04-02 18:52:12 +0000292 with self.assertRaises(ZeroDivisionError):
Guido van Rossum1a5e21e2006-02-28 21:57:43 +0000293 with lock:
Benjamin Petersonc9c0f202009-06-30 23:06:06 +0000294 self.assertTrue(locked())
Florent Xicluna41fe6152010-04-02 18:52:12 +0000295 1 / 0
296 self.assertFalse(locked())
Guido van Rossum1a5e21e2006-02-28 21:57:43 +0000297
298 def testWithLock(self):
299 lock = threading.Lock()
300 self.boilerPlate(lock, lock.locked)
301
302 def testWithRLock(self):
303 lock = threading.RLock()
304 self.boilerPlate(lock, lock._is_owned)
305
306 def testWithCondition(self):
307 lock = threading.Condition()
308 def locked():
309 return lock._is_owned()
310 self.boilerPlate(lock, locked)
311
312 def testWithSemaphore(self):
313 lock = threading.Semaphore()
314 def locked():
315 if lock.acquire(False):
316 lock.release()
317 return False
318 else:
319 return True
320 self.boilerPlate(lock, locked)
321
322 def testWithBoundedSemaphore(self):
323 lock = threading.BoundedSemaphore()
324 def locked():
325 if lock.acquire(False):
326 lock.release()
327 return False
328 else:
329 return True
330 self.boilerPlate(lock, locked)
331
Michael Foordb3a89842010-06-30 12:17:50 +0000332
333class mycontext(ContextDecorator):
Nick Coghlan059def52013-10-26 18:08:15 +1000334 """Example decoration-compatible context manager for testing"""
Michael Foordb3a89842010-06-30 12:17:50 +0000335 started = False
336 exc = None
337 catch = False
338
339 def __enter__(self):
340 self.started = True
341 return self
342
343 def __exit__(self, *exc):
344 self.exc = exc
345 return self.catch
346
347
348class TestContextDecorator(unittest.TestCase):
349
Nick Coghlan561eb5c2013-10-26 22:20:43 +1000350 @support.requires_docstrings
Nick Coghlan059def52013-10-26 18:08:15 +1000351 def test_instance_docs(self):
352 # Issue 19330: ensure context manager instances have good docstrings
353 cm_docstring = mycontext.__doc__
354 obj = mycontext()
355 self.assertEqual(obj.__doc__, cm_docstring)
356
Michael Foordb3a89842010-06-30 12:17:50 +0000357 def test_contextdecorator(self):
358 context = mycontext()
359 with context as result:
360 self.assertIs(result, context)
361 self.assertTrue(context.started)
362
363 self.assertEqual(context.exc, (None, None, None))
364
365
366 def test_contextdecorator_with_exception(self):
367 context = mycontext()
368
Ezio Melottied3a7d22010-12-01 02:32:32 +0000369 with self.assertRaisesRegex(NameError, 'foo'):
Michael Foordb3a89842010-06-30 12:17:50 +0000370 with context:
371 raise NameError('foo')
372 self.assertIsNotNone(context.exc)
373 self.assertIs(context.exc[0], NameError)
374
375 context = mycontext()
376 context.catch = True
377 with context:
378 raise NameError('foo')
379 self.assertIsNotNone(context.exc)
380 self.assertIs(context.exc[0], NameError)
381
382
383 def test_decorator(self):
384 context = mycontext()
385
386 @context
387 def test():
388 self.assertIsNone(context.exc)
389 self.assertTrue(context.started)
390 test()
391 self.assertEqual(context.exc, (None, None, None))
392
393
394 def test_decorator_with_exception(self):
395 context = mycontext()
396
397 @context
398 def test():
399 self.assertIsNone(context.exc)
400 self.assertTrue(context.started)
401 raise NameError('foo')
402
Ezio Melottied3a7d22010-12-01 02:32:32 +0000403 with self.assertRaisesRegex(NameError, 'foo'):
Michael Foordb3a89842010-06-30 12:17:50 +0000404 test()
405 self.assertIsNotNone(context.exc)
406 self.assertIs(context.exc[0], NameError)
407
408
409 def test_decorating_method(self):
410 context = mycontext()
411
412 class Test(object):
413
414 @context
415 def method(self, a, b, c=None):
416 self.a = a
417 self.b = b
418 self.c = c
419
420 # these tests are for argument passing when used as a decorator
421 test = Test()
422 test.method(1, 2)
423 self.assertEqual(test.a, 1)
424 self.assertEqual(test.b, 2)
425 self.assertEqual(test.c, None)
426
427 test = Test()
428 test.method('a', 'b', 'c')
429 self.assertEqual(test.a, 'a')
430 self.assertEqual(test.b, 'b')
431 self.assertEqual(test.c, 'c')
432
433 test = Test()
434 test.method(a=1, b=2)
435 self.assertEqual(test.a, 1)
436 self.assertEqual(test.b, 2)
437
438
439 def test_typo_enter(self):
440 class mycontext(ContextDecorator):
441 def __unter__(self):
442 pass
443 def __exit__(self, *exc):
444 pass
445
446 with self.assertRaises(AttributeError):
447 with mycontext():
448 pass
449
450
451 def test_typo_exit(self):
452 class mycontext(ContextDecorator):
453 def __enter__(self):
454 pass
455 def __uxit__(self, *exc):
456 pass
457
458 with self.assertRaises(AttributeError):
459 with mycontext():
460 pass
461
462
463 def test_contextdecorator_as_mixin(self):
464 class somecontext(object):
465 started = False
466 exc = None
467
468 def __enter__(self):
469 self.started = True
470 return self
471
472 def __exit__(self, *exc):
473 self.exc = exc
474
475 class mycontext(somecontext, ContextDecorator):
476 pass
477
478 context = mycontext()
479 @context
480 def test():
481 self.assertIsNone(context.exc)
482 self.assertTrue(context.started)
483 test()
484 self.assertEqual(context.exc, (None, None, None))
485
486
487 def test_contextmanager_as_decorator(self):
Michael Foordb3a89842010-06-30 12:17:50 +0000488 @contextmanager
489 def woohoo(y):
490 state.append(y)
491 yield
492 state.append(999)
493
Nick Coghlan0ded3e32011-05-05 23:49:25 +1000494 state = []
Michael Foordb3a89842010-06-30 12:17:50 +0000495 @woohoo(1)
496 def test(x):
497 self.assertEqual(state, [1])
498 state.append(x)
499 test('something')
500 self.assertEqual(state, [1, 'something', 999])
501
Nick Coghlan0ded3e32011-05-05 23:49:25 +1000502 # Issue #11647: Ensure the decorated function is 'reusable'
503 state = []
504 test('something else')
505 self.assertEqual(state, [1, 'something else', 999])
506
Michael Foordb3a89842010-06-30 12:17:50 +0000507
Nick Coghlan3267a302012-05-21 22:54:43 +1000508class TestExitStack(unittest.TestCase):
509
Nick Coghlan561eb5c2013-10-26 22:20:43 +1000510 @support.requires_docstrings
Nick Coghlan059def52013-10-26 18:08:15 +1000511 def test_instance_docs(self):
512 # Issue 19330: ensure context manager instances have good docstrings
513 cm_docstring = ExitStack.__doc__
514 obj = ExitStack()
515 self.assertEqual(obj.__doc__, cm_docstring)
516
Nick Coghlan3267a302012-05-21 22:54:43 +1000517 def test_no_resources(self):
518 with ExitStack():
519 pass
520
521 def test_callback(self):
522 expected = [
523 ((), {}),
524 ((1,), {}),
525 ((1,2), {}),
526 ((), dict(example=1)),
527 ((1,), dict(example=1)),
528 ((1,2), dict(example=1)),
529 ]
530 result = []
531 def _exit(*args, **kwds):
532 """Test metadata propagation"""
533 result.append((args, kwds))
534 with ExitStack() as stack:
535 for args, kwds in reversed(expected):
536 if args and kwds:
537 f = stack.callback(_exit, *args, **kwds)
538 elif args:
539 f = stack.callback(_exit, *args)
540 elif kwds:
541 f = stack.callback(_exit, **kwds)
542 else:
543 f = stack.callback(_exit)
544 self.assertIs(f, _exit)
545 for wrapper in stack._exit_callbacks:
546 self.assertIs(wrapper.__wrapped__, _exit)
547 self.assertNotEqual(wrapper.__name__, _exit.__name__)
548 self.assertIsNone(wrapper.__doc__, _exit.__doc__)
549 self.assertEqual(result, expected)
550
551 def test_push(self):
552 exc_raised = ZeroDivisionError
553 def _expect_exc(exc_type, exc, exc_tb):
554 self.assertIs(exc_type, exc_raised)
555 def _suppress_exc(*exc_details):
556 return True
557 def _expect_ok(exc_type, exc, exc_tb):
558 self.assertIsNone(exc_type)
559 self.assertIsNone(exc)
560 self.assertIsNone(exc_tb)
561 class ExitCM(object):
562 def __init__(self, check_exc):
563 self.check_exc = check_exc
564 def __enter__(self):
565 self.fail("Should not be called!")
566 def __exit__(self, *exc_details):
567 self.check_exc(*exc_details)
568 with ExitStack() as stack:
569 stack.push(_expect_ok)
570 self.assertIs(stack._exit_callbacks[-1], _expect_ok)
571 cm = ExitCM(_expect_ok)
572 stack.push(cm)
573 self.assertIs(stack._exit_callbacks[-1].__self__, cm)
574 stack.push(_suppress_exc)
575 self.assertIs(stack._exit_callbacks[-1], _suppress_exc)
576 cm = ExitCM(_expect_exc)
577 stack.push(cm)
578 self.assertIs(stack._exit_callbacks[-1].__self__, cm)
579 stack.push(_expect_exc)
580 self.assertIs(stack._exit_callbacks[-1], _expect_exc)
581 stack.push(_expect_exc)
582 self.assertIs(stack._exit_callbacks[-1], _expect_exc)
583 1/0
584
585 def test_enter_context(self):
586 class TestCM(object):
587 def __enter__(self):
588 result.append(1)
589 def __exit__(self, *exc_details):
590 result.append(3)
591
592 result = []
593 cm = TestCM()
594 with ExitStack() as stack:
595 @stack.callback # Registered first => cleaned up last
596 def _exit():
597 result.append(4)
598 self.assertIsNotNone(_exit)
599 stack.enter_context(cm)
600 self.assertIs(stack._exit_callbacks[-1].__self__, cm)
601 result.append(2)
602 self.assertEqual(result, [1, 2, 3, 4])
603
604 def test_close(self):
605 result = []
606 with ExitStack() as stack:
607 @stack.callback
608 def _exit():
609 result.append(1)
610 self.assertIsNotNone(_exit)
611 stack.close()
612 result.append(2)
613 self.assertEqual(result, [1, 2])
614
615 def test_pop_all(self):
616 result = []
617 with ExitStack() as stack:
618 @stack.callback
619 def _exit():
620 result.append(3)
621 self.assertIsNotNone(_exit)
622 new_stack = stack.pop_all()
623 result.append(1)
624 result.append(2)
625 new_stack.close()
626 self.assertEqual(result, [1, 2, 3])
627
Nick Coghlanc73e8c22012-05-31 23:49:26 +1000628 def test_exit_raise(self):
629 with self.assertRaises(ZeroDivisionError):
630 with ExitStack() as stack:
631 stack.push(lambda *exc: False)
632 1/0
633
634 def test_exit_suppress(self):
635 with ExitStack() as stack:
636 stack.push(lambda *exc: True)
637 1/0
638
639 def test_exit_exception_chaining_reference(self):
640 # Sanity check to make sure that ExitStack chaining matches
641 # actual nested with statements
642 class RaiseExc:
643 def __init__(self, exc):
644 self.exc = exc
645 def __enter__(self):
646 return self
647 def __exit__(self, *exc_details):
648 raise self.exc
649
Nick Coghlan77452fc2012-06-01 22:48:32 +1000650 class RaiseExcWithContext:
651 def __init__(self, outer, inner):
652 self.outer = outer
653 self.inner = inner
654 def __enter__(self):
655 return self
656 def __exit__(self, *exc_details):
657 try:
658 raise self.inner
659 except:
660 raise self.outer
661
Nick Coghlanc73e8c22012-05-31 23:49:26 +1000662 class SuppressExc:
663 def __enter__(self):
664 return self
665 def __exit__(self, *exc_details):
666 type(self).saved_details = exc_details
667 return True
668
669 try:
670 with RaiseExc(IndexError):
Nick Coghlan77452fc2012-06-01 22:48:32 +1000671 with RaiseExcWithContext(KeyError, AttributeError):
672 with SuppressExc():
673 with RaiseExc(ValueError):
674 1 / 0
Nick Coghlanc73e8c22012-05-31 23:49:26 +1000675 except IndexError as exc:
676 self.assertIsInstance(exc.__context__, KeyError)
677 self.assertIsInstance(exc.__context__.__context__, AttributeError)
678 # Inner exceptions were suppressed
679 self.assertIsNone(exc.__context__.__context__.__context__)
680 else:
681 self.fail("Expected IndexError, but no exception was raised")
682 # Check the inner exceptions
683 inner_exc = SuppressExc.saved_details[1]
684 self.assertIsInstance(inner_exc, ValueError)
685 self.assertIsInstance(inner_exc.__context__, ZeroDivisionError)
686
687 def test_exit_exception_chaining(self):
688 # Ensure exception chaining matches the reference behaviour
689 def raise_exc(exc):
690 raise exc
691
692 saved_details = None
693 def suppress_exc(*exc_details):
694 nonlocal saved_details
695 saved_details = exc_details
696 return True
697
698 try:
699 with ExitStack() as stack:
700 stack.callback(raise_exc, IndexError)
701 stack.callback(raise_exc, KeyError)
702 stack.callback(raise_exc, AttributeError)
703 stack.push(suppress_exc)
704 stack.callback(raise_exc, ValueError)
705 1 / 0
706 except IndexError as exc:
707 self.assertIsInstance(exc.__context__, KeyError)
708 self.assertIsInstance(exc.__context__.__context__, AttributeError)
Nick Coghlan77452fc2012-06-01 22:48:32 +1000709 # Inner exceptions were suppressed
710 self.assertIsNone(exc.__context__.__context__.__context__)
Nick Coghlanc73e8c22012-05-31 23:49:26 +1000711 else:
712 self.fail("Expected IndexError, but no exception was raised")
713 # Check the inner exceptions
714 inner_exc = saved_details[1]
715 self.assertIsInstance(inner_exc, ValueError)
716 self.assertIsInstance(inner_exc.__context__, ZeroDivisionError)
717
Nick Coghlan1a33b2f2013-10-01 23:24:56 +1000718 def test_exit_exception_non_suppressing(self):
719 # http://bugs.python.org/issue19092
720 def raise_exc(exc):
721 raise exc
722
723 def suppress_exc(*exc_details):
724 return True
725
726 try:
727 with ExitStack() as stack:
728 stack.callback(lambda: None)
729 stack.callback(raise_exc, IndexError)
730 except Exception as exc:
731 self.assertIsInstance(exc, IndexError)
732 else:
733 self.fail("Expected IndexError, but no exception was raised")
734
735 try:
736 with ExitStack() as stack:
737 stack.callback(raise_exc, KeyError)
738 stack.push(suppress_exc)
739 stack.callback(raise_exc, IndexError)
740 except Exception as exc:
741 self.assertIsInstance(exc, KeyError)
742 else:
743 self.fail("Expected KeyError, but no exception was raised")
744
Nick Coghlan09761e72014-01-22 22:24:46 +1000745 def test_exit_exception_with_correct_context(self):
746 # http://bugs.python.org/issue20317
747 @contextmanager
Nick Coghlanadd94c92014-01-24 23:05:45 +1000748 def gets_the_context_right(exc):
Nick Coghlan09761e72014-01-22 22:24:46 +1000749 try:
Nick Coghlanadd94c92014-01-24 23:05:45 +1000750 yield
Nick Coghlan09761e72014-01-22 22:24:46 +1000751 finally:
Nick Coghlanadd94c92014-01-24 23:05:45 +1000752 raise exc
753
754 exc1 = Exception(1)
755 exc2 = Exception(2)
756 exc3 = Exception(3)
757 exc4 = Exception(4)
Nick Coghlan09761e72014-01-22 22:24:46 +1000758
759 # The contextmanager already fixes the context, so prior to the
760 # fix, ExitStack would try to fix it *again* and get into an
761 # infinite self-referential loop
762 try:
763 with ExitStack() as stack:
Nick Coghlanadd94c92014-01-24 23:05:45 +1000764 stack.enter_context(gets_the_context_right(exc4))
765 stack.enter_context(gets_the_context_right(exc3))
766 stack.enter_context(gets_the_context_right(exc2))
767 raise exc1
768 except Exception as exc:
769 self.assertIs(exc, exc4)
770 self.assertIs(exc.__context__, exc3)
771 self.assertIs(exc.__context__.__context__, exc2)
772 self.assertIs(exc.__context__.__context__.__context__, exc1)
773 self.assertIsNone(
774 exc.__context__.__context__.__context__.__context__)
775
776 def test_exit_exception_with_existing_context(self):
777 # Addresses a lack of test coverage discovered after checking in a
778 # fix for issue 20317 that still contained debugging code.
779 def raise_nested(inner_exc, outer_exc):
780 try:
781 raise inner_exc
782 finally:
783 raise outer_exc
784 exc1 = Exception(1)
785 exc2 = Exception(2)
786 exc3 = Exception(3)
787 exc4 = Exception(4)
788 exc5 = Exception(5)
789 try:
790 with ExitStack() as stack:
791 stack.callback(raise_nested, exc4, exc5)
792 stack.callback(raise_nested, exc2, exc3)
793 raise exc1
794 except Exception as exc:
795 self.assertIs(exc, exc5)
796 self.assertIs(exc.__context__, exc4)
797 self.assertIs(exc.__context__.__context__, exc3)
798 self.assertIs(exc.__context__.__context__.__context__, exc2)
799 self.assertIs(
800 exc.__context__.__context__.__context__.__context__, exc1)
801 self.assertIsNone(
802 exc.__context__.__context__.__context__.__context__.__context__)
803
Nick Coghlan09761e72014-01-22 22:24:46 +1000804
805
Nick Coghlan1a33b2f2013-10-01 23:24:56 +1000806 def test_body_exception_suppress(self):
807 def suppress_exc(*exc_details):
808 return True
809 try:
810 with ExitStack() as stack:
811 stack.push(suppress_exc)
812 1/0
813 except IndexError as exc:
814 self.fail("Expected no exception, got IndexError")
815
Nick Coghlanc73e8c22012-05-31 23:49:26 +1000816 def test_exit_exception_chaining_suppress(self):
817 with ExitStack() as stack:
818 stack.push(lambda *exc: True)
819 stack.push(lambda *exc: 1/0)
820 stack.push(lambda *exc: {}[1])
821
Nick Coghlana5bd2a12012-06-01 00:00:38 +1000822 def test_excessive_nesting(self):
823 # The original implementation would die with RecursionError here
824 with ExitStack() as stack:
825 for i in range(10000):
826 stack.callback(int)
827
Nick Coghlan3267a302012-05-21 22:54:43 +1000828 def test_instance_bypass(self):
829 class Example(object): pass
830 cm = Example()
831 cm.__exit__ = object()
832 stack = ExitStack()
833 self.assertRaises(AttributeError, stack.enter_context, cm)
834 stack.push(cm)
835 self.assertIs(stack._exit_callbacks[-1], cm)
836
Gregory P. Smithba2ecd62016-06-14 09:19:20 -0700837 def test_dont_reraise_RuntimeError(self):
Serhiy Storchakace1a9f32016-06-20 05:29:54 +0300838 # https://bugs.python.org/issue27122
Gregory P. Smithba2ecd62016-06-14 09:19:20 -0700839 class UniqueException(Exception): pass
Serhiy Storchakace1a9f32016-06-20 05:29:54 +0300840 class UniqueRuntimeError(RuntimeError): pass
Gregory P. Smithba2ecd62016-06-14 09:19:20 -0700841
842 @contextmanager
843 def second():
844 try:
845 yield 1
846 except Exception as exc:
847 raise UniqueException("new exception") from exc
848
849 @contextmanager
850 def first():
851 try:
852 yield 1
853 except Exception as exc:
854 raise exc
855
Serhiy Storchakace1a9f32016-06-20 05:29:54 +0300856 # The UniqueRuntimeError should be caught by second()'s exception
Gregory P. Smithba2ecd62016-06-14 09:19:20 -0700857 # handler which chain raised a new UniqueException.
858 with self.assertRaises(UniqueException) as err_ctx:
859 with ExitStack() as es_ctx:
860 es_ctx.enter_context(second())
861 es_ctx.enter_context(first())
Serhiy Storchakace1a9f32016-06-20 05:29:54 +0300862 raise UniqueRuntimeError("please no infinite loop.")
Gregory P. Smithba2ecd62016-06-14 09:19:20 -0700863
Serhiy Storchakace1a9f32016-06-20 05:29:54 +0300864 exc = err_ctx.exception
865 self.assertIsInstance(exc, UniqueException)
866 self.assertIsInstance(exc.__context__, UniqueRuntimeError)
867 self.assertIsNone(exc.__context__.__context__)
868 self.assertIsNone(exc.__context__.__cause__)
869 self.assertIs(exc.__cause__, exc.__context__)
Gregory P. Smithba2ecd62016-06-14 09:19:20 -0700870
Berker Peksagbb44fe02014-11-28 23:28:06 +0200871
872class TestRedirectStream:
873
874 redirect_stream = None
875 orig_stream = None
Raymond Hettinger088cbf22013-10-10 00:46:57 -0700876
Nick Coghlan561eb5c2013-10-26 22:20:43 +1000877 @support.requires_docstrings
Nick Coghlan059def52013-10-26 18:08:15 +1000878 def test_instance_docs(self):
879 # Issue 19330: ensure context manager instances have good docstrings
Berker Peksagbb44fe02014-11-28 23:28:06 +0200880 cm_docstring = self.redirect_stream.__doc__
881 obj = self.redirect_stream(None)
Nick Coghlan059def52013-10-26 18:08:15 +1000882 self.assertEqual(obj.__doc__, cm_docstring)
883
Nick Coghlan8e113b42013-11-03 17:00:51 +1000884 def test_no_redirect_in_init(self):
Berker Peksagbb44fe02014-11-28 23:28:06 +0200885 orig_stdout = getattr(sys, self.orig_stream)
886 self.redirect_stream(None)
887 self.assertIs(getattr(sys, self.orig_stream), orig_stdout)
Nick Coghlan8e113b42013-11-03 17:00:51 +1000888
Raymond Hettinger088cbf22013-10-10 00:46:57 -0700889 def test_redirect_to_string_io(self):
890 f = io.StringIO()
Nick Coghlan0ddaed32013-10-26 16:37:47 +1000891 msg = "Consider an API like help(), which prints directly to stdout"
Berker Peksagbb44fe02014-11-28 23:28:06 +0200892 orig_stdout = getattr(sys, self.orig_stream)
893 with self.redirect_stream(f):
894 print(msg, file=getattr(sys, self.orig_stream))
895 self.assertIs(getattr(sys, self.orig_stream), orig_stdout)
Nick Coghlan0ddaed32013-10-26 16:37:47 +1000896 s = f.getvalue().strip()
897 self.assertEqual(s, msg)
Nick Coghlan3267a302012-05-21 22:54:43 +1000898
Nick Coghlan8608d262013-10-20 00:30:51 +1000899 def test_enter_result_is_target(self):
900 f = io.StringIO()
Berker Peksagbb44fe02014-11-28 23:28:06 +0200901 with self.redirect_stream(f) as enter_result:
Nick Coghlan8608d262013-10-20 00:30:51 +1000902 self.assertIs(enter_result, f)
903
904 def test_cm_is_reusable(self):
905 f = io.StringIO()
Berker Peksagbb44fe02014-11-28 23:28:06 +0200906 write_to_f = self.redirect_stream(f)
907 orig_stdout = getattr(sys, self.orig_stream)
Nick Coghlan8608d262013-10-20 00:30:51 +1000908 with write_to_f:
Berker Peksagbb44fe02014-11-28 23:28:06 +0200909 print("Hello", end=" ", file=getattr(sys, self.orig_stream))
Nick Coghlan8608d262013-10-20 00:30:51 +1000910 with write_to_f:
Berker Peksagbb44fe02014-11-28 23:28:06 +0200911 print("World!", file=getattr(sys, self.orig_stream))
912 self.assertIs(getattr(sys, self.orig_stream), orig_stdout)
Nick Coghlan8608d262013-10-20 00:30:51 +1000913 s = f.getvalue()
914 self.assertEqual(s, "Hello World!\n")
915
Nick Coghlan8e113b42013-11-03 17:00:51 +1000916 def test_cm_is_reentrant(self):
Nick Coghlan8608d262013-10-20 00:30:51 +1000917 f = io.StringIO()
Berker Peksagbb44fe02014-11-28 23:28:06 +0200918 write_to_f = self.redirect_stream(f)
919 orig_stdout = getattr(sys, self.orig_stream)
Nick Coghlan8e113b42013-11-03 17:00:51 +1000920 with write_to_f:
Berker Peksagbb44fe02014-11-28 23:28:06 +0200921 print("Hello", end=" ", file=getattr(sys, self.orig_stream))
Nick Coghlan8608d262013-10-20 00:30:51 +1000922 with write_to_f:
Berker Peksagbb44fe02014-11-28 23:28:06 +0200923 print("World!", file=getattr(sys, self.orig_stream))
924 self.assertIs(getattr(sys, self.orig_stream), orig_stdout)
Nick Coghlan8e113b42013-11-03 17:00:51 +1000925 s = f.getvalue()
926 self.assertEqual(s, "Hello World!\n")
Nick Coghlan8608d262013-10-20 00:30:51 +1000927
928
Berker Peksagbb44fe02014-11-28 23:28:06 +0200929class TestRedirectStdout(TestRedirectStream, unittest.TestCase):
930
931 redirect_stream = redirect_stdout
932 orig_stream = "stdout"
933
934
935class TestRedirectStderr(TestRedirectStream, unittest.TestCase):
936
937 redirect_stream = redirect_stderr
938 orig_stream = "stderr"
939
940
Nick Coghlan240f86d2013-10-17 23:40:57 +1000941class TestSuppress(unittest.TestCase):
942
Nick Coghlan561eb5c2013-10-26 22:20:43 +1000943 @support.requires_docstrings
Nick Coghlan059def52013-10-26 18:08:15 +1000944 def test_instance_docs(self):
945 # Issue 19330: ensure context manager instances have good docstrings
946 cm_docstring = suppress.__doc__
947 obj = suppress()
948 self.assertEqual(obj.__doc__, cm_docstring)
949
Nick Coghlan8608d262013-10-20 00:30:51 +1000950 def test_no_result_from_enter(self):
951 with suppress(ValueError) as enter_result:
952 self.assertIsNone(enter_result)
Nick Coghlan240f86d2013-10-17 23:40:57 +1000953
Nick Coghlan8608d262013-10-20 00:30:51 +1000954 def test_no_exception(self):
Nick Coghlan240f86d2013-10-17 23:40:57 +1000955 with suppress(ValueError):
956 self.assertEqual(pow(2, 5), 32)
957
958 def test_exact_exception(self):
Nick Coghlan240f86d2013-10-17 23:40:57 +1000959 with suppress(TypeError):
960 len(5)
961
Nick Coghlan059def52013-10-26 18:08:15 +1000962 def test_exception_hierarchy(self):
963 with suppress(LookupError):
964 'Hello'[50]
965
966 def test_other_exception(self):
967 with self.assertRaises(ZeroDivisionError):
968 with suppress(TypeError):
969 1/0
970
971 def test_no_args(self):
972 with self.assertRaises(ZeroDivisionError):
973 with suppress():
974 1/0
975
Nick Coghlan240f86d2013-10-17 23:40:57 +1000976 def test_multiple_exception_args(self):
Nick Coghlan8608d262013-10-20 00:30:51 +1000977 with suppress(ZeroDivisionError, TypeError):
978 1/0
Nick Coghlan240f86d2013-10-17 23:40:57 +1000979 with suppress(ZeroDivisionError, TypeError):
980 len(5)
981
Nick Coghlan8608d262013-10-20 00:30:51 +1000982 def test_cm_is_reentrant(self):
983 ignore_exceptions = suppress(Exception)
984 with ignore_exceptions:
985 pass
986 with ignore_exceptions:
987 len(5)
988 with ignore_exceptions:
Nick Coghlan8608d262013-10-20 00:30:51 +1000989 with ignore_exceptions: # Check nested usage
990 len(5)
Martin Panter7c6420a2015-10-10 11:04:44 +0000991 outer_continued = True
992 1/0
993 self.assertTrue(outer_continued)
Nick Coghlan8608d262013-10-20 00:30:51 +1000994
Guido van Rossum1a5e21e2006-02-28 21:57:43 +0000995if __name__ == "__main__":
Brett Cannon3e9a9ae2013-06-12 21:25:59 -0400996 unittest.main()