blob: 188a29d9f9fd77ec5c79d1064828b81431b70910 [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
Martin Teichmanndd0e0872018-01-28 05:17:46 +010010import weakref
Guido van Rossum1a5e21e2006-02-28 21:57:43 +000011
Florent Xicluna41fe6152010-04-02 18:52:12 +000012
Brett Cannon9e080e02016-04-08 12:15:27 -070013class TestAbstractContextManager(unittest.TestCase):
14
15 def test_enter(self):
16 class DefaultEnter(AbstractContextManager):
17 def __exit__(self, *args):
18 super().__exit__(*args)
19
20 manager = DefaultEnter()
21 self.assertIs(manager.__enter__(), manager)
22
23 def test_exit_is_abstract(self):
24 class MissingExit(AbstractContextManager):
25 pass
26
27 with self.assertRaises(TypeError):
28 MissingExit()
29
30 def test_structural_subclassing(self):
31 class ManagerFromScratch:
32 def __enter__(self):
33 return self
34 def __exit__(self, exc_type, exc_value, traceback):
35 return None
36
37 self.assertTrue(issubclass(ManagerFromScratch, AbstractContextManager))
38
39 class DefaultEnter(AbstractContextManager):
40 def __exit__(self, *args):
41 super().__exit__(*args)
42
43 self.assertTrue(issubclass(DefaultEnter, AbstractContextManager))
44
Jelle Zijlstra57161aa2017-06-09 08:21:47 -070045 class NoEnter(ManagerFromScratch):
46 __enter__ = None
47
48 self.assertFalse(issubclass(NoEnter, AbstractContextManager))
49
50 class NoExit(ManagerFromScratch):
51 __exit__ = None
52
53 self.assertFalse(issubclass(NoExit, AbstractContextManager))
54
Brett Cannon9e080e02016-04-08 12:15:27 -070055
Guido van Rossum1a5e21e2006-02-28 21:57:43 +000056class ContextManagerTestCase(unittest.TestCase):
57
58 def test_contextmanager_plain(self):
59 state = []
60 @contextmanager
61 def woohoo():
62 state.append(1)
63 yield 42
64 state.append(999)
65 with woohoo() as x:
66 self.assertEqual(state, [1])
67 self.assertEqual(x, 42)
68 state.append(x)
69 self.assertEqual(state, [1, 42, 999])
70
71 def test_contextmanager_finally(self):
72 state = []
73 @contextmanager
74 def woohoo():
75 state.append(1)
76 try:
77 yield 42
78 finally:
79 state.append(999)
Florent Xicluna41fe6152010-04-02 18:52:12 +000080 with self.assertRaises(ZeroDivisionError):
Guido van Rossum1a5e21e2006-02-28 21:57:43 +000081 with woohoo() as x:
82 self.assertEqual(state, [1])
83 self.assertEqual(x, 42)
84 state.append(x)
85 raise ZeroDivisionError()
Guido van Rossum1a5e21e2006-02-28 21:57:43 +000086 self.assertEqual(state, [1, 42, 999])
87
Thomas Wouters49fd7fa2006-04-21 10:40:58 +000088 def test_contextmanager_no_reraise(self):
89 @contextmanager
90 def whee():
91 yield
Thomas Wouters477c8d52006-05-27 19:21:47 +000092 ctx = whee()
Thomas Wouters49fd7fa2006-04-21 10:40:58 +000093 ctx.__enter__()
94 # Calling __exit__ should not result in an exception
Benjamin Petersonc9c0f202009-06-30 23:06:06 +000095 self.assertFalse(ctx.__exit__(TypeError, TypeError("foo"), None))
Thomas Wouters49fd7fa2006-04-21 10:40:58 +000096
97 def test_contextmanager_trap_yield_after_throw(self):
98 @contextmanager
99 def whoo():
100 try:
101 yield
102 except:
103 yield
Thomas Wouters477c8d52006-05-27 19:21:47 +0000104 ctx = whoo()
Thomas Wouters49fd7fa2006-04-21 10:40:58 +0000105 ctx.__enter__()
106 self.assertRaises(
107 RuntimeError, ctx.__exit__, TypeError, TypeError("foo"), None
108 )
109
Guido van Rossum1a5e21e2006-02-28 21:57:43 +0000110 def test_contextmanager_except(self):
111 state = []
112 @contextmanager
113 def woohoo():
114 state.append(1)
115 try:
116 yield 42
Guido van Rossumb940e112007-01-10 16:19:56 +0000117 except ZeroDivisionError as e:
Guido van Rossum1a5e21e2006-02-28 21:57:43 +0000118 state.append(e.args[0])
119 self.assertEqual(state, [1, 42, 999])
120 with woohoo() as x:
121 self.assertEqual(state, [1])
122 self.assertEqual(x, 42)
123 state.append(x)
124 raise ZeroDivisionError(999)
125 self.assertEqual(state, [1, 42, 999])
126
Yury Selivanov8170e8c2015-05-09 11:44:30 -0400127 def test_contextmanager_except_stopiter(self):
128 stop_exc = StopIteration('spam')
129 @contextmanager
130 def woohoo():
131 yield
132 try:
Martin Panter7e3a91a2016-02-10 04:40:48 +0000133 with self.assertWarnsRegex(DeprecationWarning,
Yury Selivanov68333392015-05-22 11:16:47 -0400134 "StopIteration"):
135 with woohoo():
136 raise stop_exc
Yury Selivanov8170e8c2015-05-09 11:44:30 -0400137 except Exception as ex:
138 self.assertIs(ex, stop_exc)
139 else:
140 self.fail('StopIteration was suppressed')
141
142 def test_contextmanager_except_pep479(self):
143 code = """\
144from __future__ import generator_stop
145from contextlib import contextmanager
146@contextmanager
147def woohoo():
148 yield
149"""
150 locals = {}
151 exec(code, locals, locals)
152 woohoo = locals['woohoo']
153
154 stop_exc = StopIteration('spam')
155 try:
156 with woohoo():
157 raise stop_exc
158 except Exception as ex:
159 self.assertIs(ex, stop_exc)
160 else:
161 self.fail('StopIteration was suppressed')
162
svelankar00c75e92017-04-11 05:11:13 -0400163 def test_contextmanager_do_not_unchain_non_stopiteration_exceptions(self):
164 @contextmanager
165 def test_issue29692():
166 try:
167 yield
168 except Exception as exc:
169 raise RuntimeError('issue29692:Chained') from exc
170 try:
171 with test_issue29692():
172 raise ZeroDivisionError
173 except Exception as ex:
174 self.assertIs(type(ex), RuntimeError)
175 self.assertEqual(ex.args[0], 'issue29692:Chained')
176 self.assertIsInstance(ex.__cause__, ZeroDivisionError)
177
178 try:
179 with test_issue29692():
180 raise StopIteration('issue29692:Unchained')
181 except Exception as ex:
182 self.assertIs(type(ex), StopIteration)
183 self.assertEqual(ex.args[0], 'issue29692:Unchained')
184 self.assertIsNone(ex.__cause__)
185
R. David Murray378c0cf2010-02-24 01:46:21 +0000186 def _create_contextmanager_attribs(self):
Thomas Wouters49fd7fa2006-04-21 10:40:58 +0000187 def attribs(**kw):
188 def decorate(func):
189 for k,v in kw.items():
190 setattr(func,k,v)
191 return func
192 return decorate
193 @contextmanager
194 @attribs(foo='bar')
195 def baz(spam):
196 """Whee!"""
R. David Murray378c0cf2010-02-24 01:46:21 +0000197 return baz
198
199 def test_contextmanager_attribs(self):
200 baz = self._create_contextmanager_attribs()
Thomas Wouters49fd7fa2006-04-21 10:40:58 +0000201 self.assertEqual(baz.__name__,'baz')
202 self.assertEqual(baz.foo, 'bar')
R. David Murray378c0cf2010-02-24 01:46:21 +0000203
Nick Coghlan561eb5c2013-10-26 22:20:43 +1000204 @support.requires_docstrings
R. David Murray378c0cf2010-02-24 01:46:21 +0000205 def test_contextmanager_doc_attrib(self):
206 baz = self._create_contextmanager_attribs()
Thomas Wouters49fd7fa2006-04-21 10:40:58 +0000207 self.assertEqual(baz.__doc__, "Whee!")
208
Nick Coghlan561eb5c2013-10-26 22:20:43 +1000209 @support.requires_docstrings
210 def test_instance_docstring_given_cm_docstring(self):
211 baz = self._create_contextmanager_attribs()(None)
212 self.assertEqual(baz.__doc__, "Whee!")
213
Serhiy Storchaka101ff352015-06-28 17:06:07 +0300214 def test_keywords(self):
215 # Ensure no keyword arguments are inhibited
216 @contextmanager
217 def woohoo(self, func, args, kwds):
218 yield (self, func, args, kwds)
219 with woohoo(self=11, func=22, args=33, kwds=44) as target:
220 self.assertEqual(target, (11, 22, 33, 44))
221
Martin Teichmanndd0e0872018-01-28 05:17:46 +0100222 def test_nokeepref(self):
223 class A:
224 pass
225
226 @contextmanager
227 def woohoo(a, b):
228 a = weakref.ref(a)
229 b = weakref.ref(b)
230 self.assertIsNone(a())
231 self.assertIsNone(b())
232 yield
233
234 with woohoo(A(), b=A()):
235 pass
236
237 def test_param_errors(self):
238 @contextmanager
239 def woohoo(a, *, b):
240 yield
241
242 with self.assertRaises(TypeError):
243 woohoo()
244 with self.assertRaises(TypeError):
245 woohoo(3, 5)
246 with self.assertRaises(TypeError):
247 woohoo(b=3)
248
249 def test_recursive(self):
250 depth = 0
251 @contextmanager
252 def woohoo():
253 nonlocal depth
254 before = depth
255 depth += 1
256 yield
257 depth -= 1
258 self.assertEqual(depth, before)
259
260 @woohoo()
261 def recursive():
262 if depth < 10:
263 recursive()
264
265 recursive()
266 self.assertEqual(depth, 0)
267
Nick Coghlan561eb5c2013-10-26 22:20:43 +1000268
Guido van Rossum1a5e21e2006-02-28 21:57:43 +0000269class ClosingTestCase(unittest.TestCase):
270
Nick Coghlan561eb5c2013-10-26 22:20:43 +1000271 @support.requires_docstrings
Nick Coghlan059def52013-10-26 18:08:15 +1000272 def test_instance_docs(self):
273 # Issue 19330: ensure context manager instances have good docstrings
274 cm_docstring = closing.__doc__
275 obj = closing(None)
276 self.assertEqual(obj.__doc__, cm_docstring)
Guido van Rossum1a5e21e2006-02-28 21:57:43 +0000277
278 def test_closing(self):
279 state = []
280 class C:
281 def close(self):
282 state.append(1)
283 x = C()
284 self.assertEqual(state, [])
285 with closing(x) as y:
286 self.assertEqual(x, y)
287 self.assertEqual(state, [1])
288
289 def test_closing_error(self):
290 state = []
291 class C:
292 def close(self):
293 state.append(1)
294 x = C()
295 self.assertEqual(state, [])
Florent Xicluna41fe6152010-04-02 18:52:12 +0000296 with self.assertRaises(ZeroDivisionError):
Guido van Rossum1a5e21e2006-02-28 21:57:43 +0000297 with closing(x) as y:
298 self.assertEqual(x, y)
Florent Xicluna41fe6152010-04-02 18:52:12 +0000299 1 / 0
300 self.assertEqual(state, [1])
Guido van Rossum1a5e21e2006-02-28 21:57:43 +0000301
Jesse-Bakker0784a2e2017-11-23 01:23:28 +0100302
303class NullcontextTestCase(unittest.TestCase):
304 def test_nullcontext(self):
305 class C:
306 pass
307 c = C()
308 with nullcontext(c) as c_in:
309 self.assertIs(c_in, c)
310
311
Guido van Rossum1a5e21e2006-02-28 21:57:43 +0000312class FileContextTestCase(unittest.TestCase):
313
314 def testWithOpen(self):
315 tfn = tempfile.mktemp()
316 try:
317 f = None
318 with open(tfn, "w") as f:
Benjamin Petersonc9c0f202009-06-30 23:06:06 +0000319 self.assertFalse(f.closed)
Guido van Rossum1a5e21e2006-02-28 21:57:43 +0000320 f.write("Booh\n")
Benjamin Petersonc9c0f202009-06-30 23:06:06 +0000321 self.assertTrue(f.closed)
Guido van Rossum1a5e21e2006-02-28 21:57:43 +0000322 f = None
Florent Xicluna41fe6152010-04-02 18:52:12 +0000323 with self.assertRaises(ZeroDivisionError):
Guido van Rossum1a5e21e2006-02-28 21:57:43 +0000324 with open(tfn, "r") as f:
Benjamin Petersonc9c0f202009-06-30 23:06:06 +0000325 self.assertFalse(f.closed)
Guido van Rossum1a5e21e2006-02-28 21:57:43 +0000326 self.assertEqual(f.read(), "Booh\n")
Florent Xicluna41fe6152010-04-02 18:52:12 +0000327 1 / 0
328 self.assertTrue(f.closed)
Guido van Rossum1a5e21e2006-02-28 21:57:43 +0000329 finally:
Florent Xicluna41fe6152010-04-02 18:52:12 +0000330 support.unlink(tfn)
Guido van Rossum1a5e21e2006-02-28 21:57:43 +0000331
332class LockContextTestCase(unittest.TestCase):
333
334 def boilerPlate(self, lock, locked):
Benjamin Petersonc9c0f202009-06-30 23:06:06 +0000335 self.assertFalse(locked())
Guido van Rossum1a5e21e2006-02-28 21:57:43 +0000336 with lock:
Benjamin Petersonc9c0f202009-06-30 23:06:06 +0000337 self.assertTrue(locked())
338 self.assertFalse(locked())
Florent Xicluna41fe6152010-04-02 18:52:12 +0000339 with self.assertRaises(ZeroDivisionError):
Guido van Rossum1a5e21e2006-02-28 21:57:43 +0000340 with lock:
Benjamin Petersonc9c0f202009-06-30 23:06:06 +0000341 self.assertTrue(locked())
Florent Xicluna41fe6152010-04-02 18:52:12 +0000342 1 / 0
343 self.assertFalse(locked())
Guido van Rossum1a5e21e2006-02-28 21:57:43 +0000344
345 def testWithLock(self):
346 lock = threading.Lock()
347 self.boilerPlate(lock, lock.locked)
348
349 def testWithRLock(self):
350 lock = threading.RLock()
351 self.boilerPlate(lock, lock._is_owned)
352
353 def testWithCondition(self):
354 lock = threading.Condition()
355 def locked():
356 return lock._is_owned()
357 self.boilerPlate(lock, locked)
358
359 def testWithSemaphore(self):
360 lock = threading.Semaphore()
361 def locked():
362 if lock.acquire(False):
363 lock.release()
364 return False
365 else:
366 return True
367 self.boilerPlate(lock, locked)
368
369 def testWithBoundedSemaphore(self):
370 lock = threading.BoundedSemaphore()
371 def locked():
372 if lock.acquire(False):
373 lock.release()
374 return False
375 else:
376 return True
377 self.boilerPlate(lock, locked)
378
Michael Foordb3a89842010-06-30 12:17:50 +0000379
380class mycontext(ContextDecorator):
Nick Coghlan059def52013-10-26 18:08:15 +1000381 """Example decoration-compatible context manager for testing"""
Michael Foordb3a89842010-06-30 12:17:50 +0000382 started = False
383 exc = None
384 catch = False
385
386 def __enter__(self):
387 self.started = True
388 return self
389
390 def __exit__(self, *exc):
391 self.exc = exc
392 return self.catch
393
394
395class TestContextDecorator(unittest.TestCase):
396
Nick Coghlan561eb5c2013-10-26 22:20:43 +1000397 @support.requires_docstrings
Nick Coghlan059def52013-10-26 18:08:15 +1000398 def test_instance_docs(self):
399 # Issue 19330: ensure context manager instances have good docstrings
400 cm_docstring = mycontext.__doc__
401 obj = mycontext()
402 self.assertEqual(obj.__doc__, cm_docstring)
403
Michael Foordb3a89842010-06-30 12:17:50 +0000404 def test_contextdecorator(self):
405 context = mycontext()
406 with context as result:
407 self.assertIs(result, context)
408 self.assertTrue(context.started)
409
410 self.assertEqual(context.exc, (None, None, None))
411
412
413 def test_contextdecorator_with_exception(self):
414 context = mycontext()
415
Ezio Melottied3a7d22010-12-01 02:32:32 +0000416 with self.assertRaisesRegex(NameError, 'foo'):
Michael Foordb3a89842010-06-30 12:17:50 +0000417 with context:
418 raise NameError('foo')
419 self.assertIsNotNone(context.exc)
420 self.assertIs(context.exc[0], NameError)
421
422 context = mycontext()
423 context.catch = True
424 with context:
425 raise NameError('foo')
426 self.assertIsNotNone(context.exc)
427 self.assertIs(context.exc[0], NameError)
428
429
430 def test_decorator(self):
431 context = mycontext()
432
433 @context
434 def test():
435 self.assertIsNone(context.exc)
436 self.assertTrue(context.started)
437 test()
438 self.assertEqual(context.exc, (None, None, None))
439
440
441 def test_decorator_with_exception(self):
442 context = mycontext()
443
444 @context
445 def test():
446 self.assertIsNone(context.exc)
447 self.assertTrue(context.started)
448 raise NameError('foo')
449
Ezio Melottied3a7d22010-12-01 02:32:32 +0000450 with self.assertRaisesRegex(NameError, 'foo'):
Michael Foordb3a89842010-06-30 12:17:50 +0000451 test()
452 self.assertIsNotNone(context.exc)
453 self.assertIs(context.exc[0], NameError)
454
455
456 def test_decorating_method(self):
457 context = mycontext()
458
459 class Test(object):
460
461 @context
462 def method(self, a, b, c=None):
463 self.a = a
464 self.b = b
465 self.c = c
466
467 # these tests are for argument passing when used as a decorator
468 test = Test()
469 test.method(1, 2)
470 self.assertEqual(test.a, 1)
471 self.assertEqual(test.b, 2)
472 self.assertEqual(test.c, None)
473
474 test = Test()
475 test.method('a', 'b', 'c')
476 self.assertEqual(test.a, 'a')
477 self.assertEqual(test.b, 'b')
478 self.assertEqual(test.c, 'c')
479
480 test = Test()
481 test.method(a=1, b=2)
482 self.assertEqual(test.a, 1)
483 self.assertEqual(test.b, 2)
484
485
486 def test_typo_enter(self):
487 class mycontext(ContextDecorator):
488 def __unter__(self):
489 pass
490 def __exit__(self, *exc):
491 pass
492
493 with self.assertRaises(AttributeError):
494 with mycontext():
495 pass
496
497
498 def test_typo_exit(self):
499 class mycontext(ContextDecorator):
500 def __enter__(self):
501 pass
502 def __uxit__(self, *exc):
503 pass
504
505 with self.assertRaises(AttributeError):
506 with mycontext():
507 pass
508
509
510 def test_contextdecorator_as_mixin(self):
511 class somecontext(object):
512 started = False
513 exc = None
514
515 def __enter__(self):
516 self.started = True
517 return self
518
519 def __exit__(self, *exc):
520 self.exc = exc
521
522 class mycontext(somecontext, ContextDecorator):
523 pass
524
525 context = mycontext()
526 @context
527 def test():
528 self.assertIsNone(context.exc)
529 self.assertTrue(context.started)
530 test()
531 self.assertEqual(context.exc, (None, None, None))
532
533
534 def test_contextmanager_as_decorator(self):
Michael Foordb3a89842010-06-30 12:17:50 +0000535 @contextmanager
536 def woohoo(y):
537 state.append(y)
538 yield
539 state.append(999)
540
Nick Coghlan0ded3e32011-05-05 23:49:25 +1000541 state = []
Michael Foordb3a89842010-06-30 12:17:50 +0000542 @woohoo(1)
543 def test(x):
544 self.assertEqual(state, [1])
545 state.append(x)
546 test('something')
547 self.assertEqual(state, [1, 'something', 999])
548
Nick Coghlan0ded3e32011-05-05 23:49:25 +1000549 # Issue #11647: Ensure the decorated function is 'reusable'
550 state = []
551 test('something else')
552 self.assertEqual(state, [1, 'something else', 999])
553
Michael Foordb3a89842010-06-30 12:17:50 +0000554
Ilya Kulakov1aa094f2018-01-25 12:51:18 -0800555class TestBaseExitStack:
556 exit_stack = None
Nick Coghlan3267a302012-05-21 22:54:43 +1000557
Nick Coghlan561eb5c2013-10-26 22:20:43 +1000558 @support.requires_docstrings
Nick Coghlan059def52013-10-26 18:08:15 +1000559 def test_instance_docs(self):
560 # Issue 19330: ensure context manager instances have good docstrings
Ilya Kulakov1aa094f2018-01-25 12:51:18 -0800561 cm_docstring = self.exit_stack.__doc__
562 obj = self.exit_stack()
Nick Coghlan059def52013-10-26 18:08:15 +1000563 self.assertEqual(obj.__doc__, cm_docstring)
564
Nick Coghlan3267a302012-05-21 22:54:43 +1000565 def test_no_resources(self):
Ilya Kulakov1aa094f2018-01-25 12:51:18 -0800566 with self.exit_stack():
Nick Coghlan3267a302012-05-21 22:54:43 +1000567 pass
568
569 def test_callback(self):
570 expected = [
571 ((), {}),
572 ((1,), {}),
573 ((1,2), {}),
574 ((), dict(example=1)),
575 ((1,), dict(example=1)),
576 ((1,2), dict(example=1)),
Serhiy Storchaka42a139e2019-04-01 09:16:35 +0300577 ((1,2), dict(self=3, callback=4)),
Nick Coghlan3267a302012-05-21 22:54:43 +1000578 ]
579 result = []
580 def _exit(*args, **kwds):
581 """Test metadata propagation"""
582 result.append((args, kwds))
Ilya Kulakov1aa094f2018-01-25 12:51:18 -0800583 with self.exit_stack() as stack:
Nick Coghlan3267a302012-05-21 22:54:43 +1000584 for args, kwds in reversed(expected):
585 if args and kwds:
586 f = stack.callback(_exit, *args, **kwds)
587 elif args:
588 f = stack.callback(_exit, *args)
589 elif kwds:
590 f = stack.callback(_exit, **kwds)
591 else:
592 f = stack.callback(_exit)
593 self.assertIs(f, _exit)
594 for wrapper in stack._exit_callbacks:
Ilya Kulakov1aa094f2018-01-25 12:51:18 -0800595 self.assertIs(wrapper[1].__wrapped__, _exit)
596 self.assertNotEqual(wrapper[1].__name__, _exit.__name__)
597 self.assertIsNone(wrapper[1].__doc__, _exit.__doc__)
Nick Coghlan3267a302012-05-21 22:54:43 +1000598 self.assertEqual(result, expected)
599
Serhiy Storchaka42a139e2019-04-01 09:16:35 +0300600 result = []
601 with self.exit_stack() as stack:
602 with self.assertRaises(TypeError):
603 stack.callback(arg=1)
604 with self.assertRaises(TypeError):
605 self.exit_stack.callback(arg=2)
606 with self.assertWarns(DeprecationWarning):
607 stack.callback(callback=_exit, arg=3)
608 self.assertEqual(result, [((), {'arg': 3})])
609
Nick Coghlan3267a302012-05-21 22:54:43 +1000610 def test_push(self):
611 exc_raised = ZeroDivisionError
612 def _expect_exc(exc_type, exc, exc_tb):
613 self.assertIs(exc_type, exc_raised)
614 def _suppress_exc(*exc_details):
615 return True
616 def _expect_ok(exc_type, exc, exc_tb):
617 self.assertIsNone(exc_type)
618 self.assertIsNone(exc)
619 self.assertIsNone(exc_tb)
620 class ExitCM(object):
621 def __init__(self, check_exc):
622 self.check_exc = check_exc
623 def __enter__(self):
624 self.fail("Should not be called!")
625 def __exit__(self, *exc_details):
626 self.check_exc(*exc_details)
Ilya Kulakov1aa094f2018-01-25 12:51:18 -0800627 with self.exit_stack() as stack:
Nick Coghlan3267a302012-05-21 22:54:43 +1000628 stack.push(_expect_ok)
Ilya Kulakov1aa094f2018-01-25 12:51:18 -0800629 self.assertIs(stack._exit_callbacks[-1][1], _expect_ok)
Nick Coghlan3267a302012-05-21 22:54:43 +1000630 cm = ExitCM(_expect_ok)
631 stack.push(cm)
Ilya Kulakov1aa094f2018-01-25 12:51:18 -0800632 self.assertIs(stack._exit_callbacks[-1][1].__self__, cm)
Nick Coghlan3267a302012-05-21 22:54:43 +1000633 stack.push(_suppress_exc)
Ilya Kulakov1aa094f2018-01-25 12:51:18 -0800634 self.assertIs(stack._exit_callbacks[-1][1], _suppress_exc)
Nick Coghlan3267a302012-05-21 22:54:43 +1000635 cm = ExitCM(_expect_exc)
636 stack.push(cm)
Ilya Kulakov1aa094f2018-01-25 12:51:18 -0800637 self.assertIs(stack._exit_callbacks[-1][1].__self__, cm)
Nick Coghlan3267a302012-05-21 22:54:43 +1000638 stack.push(_expect_exc)
Ilya Kulakov1aa094f2018-01-25 12:51:18 -0800639 self.assertIs(stack._exit_callbacks[-1][1], _expect_exc)
Nick Coghlan3267a302012-05-21 22:54:43 +1000640 stack.push(_expect_exc)
Ilya Kulakov1aa094f2018-01-25 12:51:18 -0800641 self.assertIs(stack._exit_callbacks[-1][1], _expect_exc)
Nick Coghlan3267a302012-05-21 22:54:43 +1000642 1/0
643
644 def test_enter_context(self):
645 class TestCM(object):
646 def __enter__(self):
647 result.append(1)
648 def __exit__(self, *exc_details):
649 result.append(3)
650
651 result = []
652 cm = TestCM()
Ilya Kulakov1aa094f2018-01-25 12:51:18 -0800653 with self.exit_stack() as stack:
Nick Coghlan3267a302012-05-21 22:54:43 +1000654 @stack.callback # Registered first => cleaned up last
655 def _exit():
656 result.append(4)
657 self.assertIsNotNone(_exit)
658 stack.enter_context(cm)
Ilya Kulakov1aa094f2018-01-25 12:51:18 -0800659 self.assertIs(stack._exit_callbacks[-1][1].__self__, cm)
Nick Coghlan3267a302012-05-21 22:54:43 +1000660 result.append(2)
661 self.assertEqual(result, [1, 2, 3, 4])
662
663 def test_close(self):
664 result = []
Ilya Kulakov1aa094f2018-01-25 12:51:18 -0800665 with self.exit_stack() as stack:
Nick Coghlan3267a302012-05-21 22:54:43 +1000666 @stack.callback
667 def _exit():
668 result.append(1)
669 self.assertIsNotNone(_exit)
670 stack.close()
671 result.append(2)
672 self.assertEqual(result, [1, 2])
673
674 def test_pop_all(self):
675 result = []
Ilya Kulakov1aa094f2018-01-25 12:51:18 -0800676 with self.exit_stack() as stack:
Nick Coghlan3267a302012-05-21 22:54:43 +1000677 @stack.callback
678 def _exit():
679 result.append(3)
680 self.assertIsNotNone(_exit)
681 new_stack = stack.pop_all()
682 result.append(1)
683 result.append(2)
684 new_stack.close()
685 self.assertEqual(result, [1, 2, 3])
686
Nick Coghlanc73e8c22012-05-31 23:49:26 +1000687 def test_exit_raise(self):
688 with self.assertRaises(ZeroDivisionError):
Ilya Kulakov1aa094f2018-01-25 12:51:18 -0800689 with self.exit_stack() as stack:
Nick Coghlanc73e8c22012-05-31 23:49:26 +1000690 stack.push(lambda *exc: False)
691 1/0
692
693 def test_exit_suppress(self):
Ilya Kulakov1aa094f2018-01-25 12:51:18 -0800694 with self.exit_stack() as stack:
Nick Coghlanc73e8c22012-05-31 23:49:26 +1000695 stack.push(lambda *exc: True)
696 1/0
697
698 def test_exit_exception_chaining_reference(self):
699 # Sanity check to make sure that ExitStack chaining matches
700 # actual nested with statements
701 class RaiseExc:
702 def __init__(self, exc):
703 self.exc = exc
704 def __enter__(self):
705 return self
706 def __exit__(self, *exc_details):
707 raise self.exc
708
Nick Coghlan77452fc2012-06-01 22:48:32 +1000709 class RaiseExcWithContext:
710 def __init__(self, outer, inner):
711 self.outer = outer
712 self.inner = inner
713 def __enter__(self):
714 return self
715 def __exit__(self, *exc_details):
716 try:
717 raise self.inner
718 except:
719 raise self.outer
720
Nick Coghlanc73e8c22012-05-31 23:49:26 +1000721 class SuppressExc:
722 def __enter__(self):
723 return self
724 def __exit__(self, *exc_details):
725 type(self).saved_details = exc_details
726 return True
727
728 try:
729 with RaiseExc(IndexError):
Nick Coghlan77452fc2012-06-01 22:48:32 +1000730 with RaiseExcWithContext(KeyError, AttributeError):
731 with SuppressExc():
732 with RaiseExc(ValueError):
733 1 / 0
Nick Coghlanc73e8c22012-05-31 23:49:26 +1000734 except IndexError as exc:
735 self.assertIsInstance(exc.__context__, KeyError)
736 self.assertIsInstance(exc.__context__.__context__, AttributeError)
737 # Inner exceptions were suppressed
738 self.assertIsNone(exc.__context__.__context__.__context__)
739 else:
740 self.fail("Expected IndexError, but no exception was raised")
741 # Check the inner exceptions
742 inner_exc = SuppressExc.saved_details[1]
743 self.assertIsInstance(inner_exc, ValueError)
744 self.assertIsInstance(inner_exc.__context__, ZeroDivisionError)
745
746 def test_exit_exception_chaining(self):
747 # Ensure exception chaining matches the reference behaviour
748 def raise_exc(exc):
749 raise exc
750
751 saved_details = None
752 def suppress_exc(*exc_details):
753 nonlocal saved_details
754 saved_details = exc_details
755 return True
756
757 try:
Ilya Kulakov1aa094f2018-01-25 12:51:18 -0800758 with self.exit_stack() as stack:
Nick Coghlanc73e8c22012-05-31 23:49:26 +1000759 stack.callback(raise_exc, IndexError)
760 stack.callback(raise_exc, KeyError)
761 stack.callback(raise_exc, AttributeError)
762 stack.push(suppress_exc)
763 stack.callback(raise_exc, ValueError)
764 1 / 0
765 except IndexError as exc:
766 self.assertIsInstance(exc.__context__, KeyError)
767 self.assertIsInstance(exc.__context__.__context__, AttributeError)
Nick Coghlan77452fc2012-06-01 22:48:32 +1000768 # Inner exceptions were suppressed
769 self.assertIsNone(exc.__context__.__context__.__context__)
Nick Coghlanc73e8c22012-05-31 23:49:26 +1000770 else:
771 self.fail("Expected IndexError, but no exception was raised")
772 # Check the inner exceptions
773 inner_exc = saved_details[1]
774 self.assertIsInstance(inner_exc, ValueError)
775 self.assertIsInstance(inner_exc.__context__, ZeroDivisionError)
776
Nick Coghlan1a33b2f2013-10-01 23:24:56 +1000777 def test_exit_exception_non_suppressing(self):
778 # http://bugs.python.org/issue19092
779 def raise_exc(exc):
780 raise exc
781
782 def suppress_exc(*exc_details):
783 return True
784
785 try:
Ilya Kulakov1aa094f2018-01-25 12:51:18 -0800786 with self.exit_stack() as stack:
Nick Coghlan1a33b2f2013-10-01 23:24:56 +1000787 stack.callback(lambda: None)
788 stack.callback(raise_exc, IndexError)
789 except Exception as exc:
790 self.assertIsInstance(exc, IndexError)
791 else:
792 self.fail("Expected IndexError, but no exception was raised")
793
794 try:
Ilya Kulakov1aa094f2018-01-25 12:51:18 -0800795 with self.exit_stack() as stack:
Nick Coghlan1a33b2f2013-10-01 23:24:56 +1000796 stack.callback(raise_exc, KeyError)
797 stack.push(suppress_exc)
798 stack.callback(raise_exc, IndexError)
799 except Exception as exc:
800 self.assertIsInstance(exc, KeyError)
801 else:
802 self.fail("Expected KeyError, but no exception was raised")
803
Nick Coghlan09761e72014-01-22 22:24:46 +1000804 def test_exit_exception_with_correct_context(self):
805 # http://bugs.python.org/issue20317
806 @contextmanager
Nick Coghlanadd94c92014-01-24 23:05:45 +1000807 def gets_the_context_right(exc):
Nick Coghlan09761e72014-01-22 22:24:46 +1000808 try:
Nick Coghlanadd94c92014-01-24 23:05:45 +1000809 yield
Nick Coghlan09761e72014-01-22 22:24:46 +1000810 finally:
Nick Coghlanadd94c92014-01-24 23:05:45 +1000811 raise exc
812
813 exc1 = Exception(1)
814 exc2 = Exception(2)
815 exc3 = Exception(3)
816 exc4 = Exception(4)
Nick Coghlan09761e72014-01-22 22:24:46 +1000817
818 # The contextmanager already fixes the context, so prior to the
819 # fix, ExitStack would try to fix it *again* and get into an
820 # infinite self-referential loop
821 try:
Ilya Kulakov1aa094f2018-01-25 12:51:18 -0800822 with self.exit_stack() as stack:
Nick Coghlanadd94c92014-01-24 23:05:45 +1000823 stack.enter_context(gets_the_context_right(exc4))
824 stack.enter_context(gets_the_context_right(exc3))
825 stack.enter_context(gets_the_context_right(exc2))
826 raise exc1
827 except Exception as exc:
828 self.assertIs(exc, exc4)
829 self.assertIs(exc.__context__, exc3)
830 self.assertIs(exc.__context__.__context__, exc2)
831 self.assertIs(exc.__context__.__context__.__context__, exc1)
832 self.assertIsNone(
833 exc.__context__.__context__.__context__.__context__)
834
835 def test_exit_exception_with_existing_context(self):
836 # Addresses a lack of test coverage discovered after checking in a
837 # fix for issue 20317 that still contained debugging code.
838 def raise_nested(inner_exc, outer_exc):
839 try:
840 raise inner_exc
841 finally:
842 raise outer_exc
843 exc1 = Exception(1)
844 exc2 = Exception(2)
845 exc3 = Exception(3)
846 exc4 = Exception(4)
847 exc5 = Exception(5)
848 try:
Ilya Kulakov1aa094f2018-01-25 12:51:18 -0800849 with self.exit_stack() as stack:
Nick Coghlanadd94c92014-01-24 23:05:45 +1000850 stack.callback(raise_nested, exc4, exc5)
851 stack.callback(raise_nested, exc2, exc3)
852 raise exc1
853 except Exception as exc:
854 self.assertIs(exc, exc5)
855 self.assertIs(exc.__context__, exc4)
856 self.assertIs(exc.__context__.__context__, exc3)
857 self.assertIs(exc.__context__.__context__.__context__, exc2)
858 self.assertIs(
859 exc.__context__.__context__.__context__.__context__, exc1)
860 self.assertIsNone(
861 exc.__context__.__context__.__context__.__context__.__context__)
862
Nick Coghlan1a33b2f2013-10-01 23:24:56 +1000863 def test_body_exception_suppress(self):
864 def suppress_exc(*exc_details):
865 return True
866 try:
Ilya Kulakov1aa094f2018-01-25 12:51:18 -0800867 with self.exit_stack() as stack:
Nick Coghlan1a33b2f2013-10-01 23:24:56 +1000868 stack.push(suppress_exc)
869 1/0
870 except IndexError as exc:
871 self.fail("Expected no exception, got IndexError")
872
Nick Coghlanc73e8c22012-05-31 23:49:26 +1000873 def test_exit_exception_chaining_suppress(self):
Ilya Kulakov1aa094f2018-01-25 12:51:18 -0800874 with self.exit_stack() as stack:
Nick Coghlanc73e8c22012-05-31 23:49:26 +1000875 stack.push(lambda *exc: True)
876 stack.push(lambda *exc: 1/0)
877 stack.push(lambda *exc: {}[1])
878
Nick Coghlana5bd2a12012-06-01 00:00:38 +1000879 def test_excessive_nesting(self):
880 # The original implementation would die with RecursionError here
Ilya Kulakov1aa094f2018-01-25 12:51:18 -0800881 with self.exit_stack() as stack:
Nick Coghlana5bd2a12012-06-01 00:00:38 +1000882 for i in range(10000):
883 stack.callback(int)
884
Nick Coghlan3267a302012-05-21 22:54:43 +1000885 def test_instance_bypass(self):
886 class Example(object): pass
887 cm = Example()
888 cm.__exit__ = object()
Ilya Kulakov1aa094f2018-01-25 12:51:18 -0800889 stack = self.exit_stack()
Nick Coghlan3267a302012-05-21 22:54:43 +1000890 self.assertRaises(AttributeError, stack.enter_context, cm)
891 stack.push(cm)
Ilya Kulakov1aa094f2018-01-25 12:51:18 -0800892 self.assertIs(stack._exit_callbacks[-1][1], cm)
Nick Coghlan3267a302012-05-21 22:54:43 +1000893
Gregory P. Smithba2ecd62016-06-14 09:19:20 -0700894 def test_dont_reraise_RuntimeError(self):
Serhiy Storchakace1a9f32016-06-20 05:29:54 +0300895 # https://bugs.python.org/issue27122
Gregory P. Smithba2ecd62016-06-14 09:19:20 -0700896 class UniqueException(Exception): pass
Serhiy Storchakace1a9f32016-06-20 05:29:54 +0300897 class UniqueRuntimeError(RuntimeError): pass
Gregory P. Smithba2ecd62016-06-14 09:19:20 -0700898
899 @contextmanager
900 def second():
901 try:
902 yield 1
903 except Exception as exc:
904 raise UniqueException("new exception") from exc
905
906 @contextmanager
907 def first():
908 try:
909 yield 1
910 except Exception as exc:
911 raise exc
912
Serhiy Storchakace1a9f32016-06-20 05:29:54 +0300913 # The UniqueRuntimeError should be caught by second()'s exception
Gregory P. Smithba2ecd62016-06-14 09:19:20 -0700914 # handler which chain raised a new UniqueException.
915 with self.assertRaises(UniqueException) as err_ctx:
Ilya Kulakov1aa094f2018-01-25 12:51:18 -0800916 with self.exit_stack() as es_ctx:
Gregory P. Smithba2ecd62016-06-14 09:19:20 -0700917 es_ctx.enter_context(second())
918 es_ctx.enter_context(first())
Serhiy Storchakace1a9f32016-06-20 05:29:54 +0300919 raise UniqueRuntimeError("please no infinite loop.")
Gregory P. Smithba2ecd62016-06-14 09:19:20 -0700920
Serhiy Storchakace1a9f32016-06-20 05:29:54 +0300921 exc = err_ctx.exception
922 self.assertIsInstance(exc, UniqueException)
923 self.assertIsInstance(exc.__context__, UniqueRuntimeError)
924 self.assertIsNone(exc.__context__.__context__)
925 self.assertIsNone(exc.__context__.__cause__)
926 self.assertIs(exc.__cause__, exc.__context__)
Gregory P. Smithba2ecd62016-06-14 09:19:20 -0700927
Berker Peksagbb44fe02014-11-28 23:28:06 +0200928
Ilya Kulakov1aa094f2018-01-25 12:51:18 -0800929class TestExitStack(TestBaseExitStack, unittest.TestCase):
930 exit_stack = ExitStack
931
932
Berker Peksagbb44fe02014-11-28 23:28:06 +0200933class TestRedirectStream:
934
935 redirect_stream = None
936 orig_stream = None
Raymond Hettinger088cbf22013-10-10 00:46:57 -0700937
Nick Coghlan561eb5c2013-10-26 22:20:43 +1000938 @support.requires_docstrings
Nick Coghlan059def52013-10-26 18:08:15 +1000939 def test_instance_docs(self):
940 # Issue 19330: ensure context manager instances have good docstrings
Berker Peksagbb44fe02014-11-28 23:28:06 +0200941 cm_docstring = self.redirect_stream.__doc__
942 obj = self.redirect_stream(None)
Nick Coghlan059def52013-10-26 18:08:15 +1000943 self.assertEqual(obj.__doc__, cm_docstring)
944
Nick Coghlan8e113b42013-11-03 17:00:51 +1000945 def test_no_redirect_in_init(self):
Berker Peksagbb44fe02014-11-28 23:28:06 +0200946 orig_stdout = getattr(sys, self.orig_stream)
947 self.redirect_stream(None)
948 self.assertIs(getattr(sys, self.orig_stream), orig_stdout)
Nick Coghlan8e113b42013-11-03 17:00:51 +1000949
Raymond Hettinger088cbf22013-10-10 00:46:57 -0700950 def test_redirect_to_string_io(self):
951 f = io.StringIO()
Nick Coghlan0ddaed32013-10-26 16:37:47 +1000952 msg = "Consider an API like help(), which prints directly to stdout"
Berker Peksagbb44fe02014-11-28 23:28:06 +0200953 orig_stdout = getattr(sys, self.orig_stream)
954 with self.redirect_stream(f):
955 print(msg, file=getattr(sys, self.orig_stream))
956 self.assertIs(getattr(sys, self.orig_stream), orig_stdout)
Nick Coghlan0ddaed32013-10-26 16:37:47 +1000957 s = f.getvalue().strip()
958 self.assertEqual(s, msg)
Nick Coghlan3267a302012-05-21 22:54:43 +1000959
Nick Coghlan8608d262013-10-20 00:30:51 +1000960 def test_enter_result_is_target(self):
961 f = io.StringIO()
Berker Peksagbb44fe02014-11-28 23:28:06 +0200962 with self.redirect_stream(f) as enter_result:
Nick Coghlan8608d262013-10-20 00:30:51 +1000963 self.assertIs(enter_result, f)
964
965 def test_cm_is_reusable(self):
966 f = io.StringIO()
Berker Peksagbb44fe02014-11-28 23:28:06 +0200967 write_to_f = self.redirect_stream(f)
968 orig_stdout = getattr(sys, self.orig_stream)
Nick Coghlan8608d262013-10-20 00:30:51 +1000969 with write_to_f:
Berker Peksagbb44fe02014-11-28 23:28:06 +0200970 print("Hello", end=" ", file=getattr(sys, self.orig_stream))
Nick Coghlan8608d262013-10-20 00:30:51 +1000971 with write_to_f:
Berker Peksagbb44fe02014-11-28 23:28:06 +0200972 print("World!", file=getattr(sys, self.orig_stream))
973 self.assertIs(getattr(sys, self.orig_stream), orig_stdout)
Nick Coghlan8608d262013-10-20 00:30:51 +1000974 s = f.getvalue()
975 self.assertEqual(s, "Hello World!\n")
976
Nick Coghlan8e113b42013-11-03 17:00:51 +1000977 def test_cm_is_reentrant(self):
Nick Coghlan8608d262013-10-20 00:30:51 +1000978 f = io.StringIO()
Berker Peksagbb44fe02014-11-28 23:28:06 +0200979 write_to_f = self.redirect_stream(f)
980 orig_stdout = getattr(sys, self.orig_stream)
Nick Coghlan8e113b42013-11-03 17:00:51 +1000981 with write_to_f:
Berker Peksagbb44fe02014-11-28 23:28:06 +0200982 print("Hello", end=" ", file=getattr(sys, self.orig_stream))
Nick Coghlan8608d262013-10-20 00:30:51 +1000983 with write_to_f:
Berker Peksagbb44fe02014-11-28 23:28:06 +0200984 print("World!", file=getattr(sys, self.orig_stream))
985 self.assertIs(getattr(sys, self.orig_stream), orig_stdout)
Nick Coghlan8e113b42013-11-03 17:00:51 +1000986 s = f.getvalue()
987 self.assertEqual(s, "Hello World!\n")
Nick Coghlan8608d262013-10-20 00:30:51 +1000988
989
Berker Peksagbb44fe02014-11-28 23:28:06 +0200990class TestRedirectStdout(TestRedirectStream, unittest.TestCase):
991
992 redirect_stream = redirect_stdout
993 orig_stream = "stdout"
994
995
996class TestRedirectStderr(TestRedirectStream, unittest.TestCase):
997
998 redirect_stream = redirect_stderr
999 orig_stream = "stderr"
1000
1001
Nick Coghlan240f86d2013-10-17 23:40:57 +10001002class TestSuppress(unittest.TestCase):
1003
Nick Coghlan561eb5c2013-10-26 22:20:43 +10001004 @support.requires_docstrings
Nick Coghlan059def52013-10-26 18:08:15 +10001005 def test_instance_docs(self):
1006 # Issue 19330: ensure context manager instances have good docstrings
1007 cm_docstring = suppress.__doc__
1008 obj = suppress()
1009 self.assertEqual(obj.__doc__, cm_docstring)
1010
Nick Coghlan8608d262013-10-20 00:30:51 +10001011 def test_no_result_from_enter(self):
1012 with suppress(ValueError) as enter_result:
1013 self.assertIsNone(enter_result)
Nick Coghlan240f86d2013-10-17 23:40:57 +10001014
Nick Coghlan8608d262013-10-20 00:30:51 +10001015 def test_no_exception(self):
Nick Coghlan240f86d2013-10-17 23:40:57 +10001016 with suppress(ValueError):
1017 self.assertEqual(pow(2, 5), 32)
1018
1019 def test_exact_exception(self):
Nick Coghlan240f86d2013-10-17 23:40:57 +10001020 with suppress(TypeError):
1021 len(5)
1022
Nick Coghlan059def52013-10-26 18:08:15 +10001023 def test_exception_hierarchy(self):
1024 with suppress(LookupError):
1025 'Hello'[50]
1026
1027 def test_other_exception(self):
1028 with self.assertRaises(ZeroDivisionError):
1029 with suppress(TypeError):
1030 1/0
1031
1032 def test_no_args(self):
1033 with self.assertRaises(ZeroDivisionError):
1034 with suppress():
1035 1/0
1036
Nick Coghlan240f86d2013-10-17 23:40:57 +10001037 def test_multiple_exception_args(self):
Nick Coghlan8608d262013-10-20 00:30:51 +10001038 with suppress(ZeroDivisionError, TypeError):
1039 1/0
Nick Coghlan240f86d2013-10-17 23:40:57 +10001040 with suppress(ZeroDivisionError, TypeError):
1041 len(5)
1042
Nick Coghlan8608d262013-10-20 00:30:51 +10001043 def test_cm_is_reentrant(self):
1044 ignore_exceptions = suppress(Exception)
1045 with ignore_exceptions:
1046 pass
1047 with ignore_exceptions:
1048 len(5)
1049 with ignore_exceptions:
Nick Coghlan8608d262013-10-20 00:30:51 +10001050 with ignore_exceptions: # Check nested usage
1051 len(5)
Martin Panter7c6420a2015-10-10 11:04:44 +00001052 outer_continued = True
1053 1/0
1054 self.assertTrue(outer_continued)
Nick Coghlan8608d262013-10-20 00:30:51 +10001055
Guido van Rossum1a5e21e2006-02-28 21:57:43 +00001056if __name__ == "__main__":
Brett Cannon3e9a9ae2013-06-12 21:25:59 -04001057 unittest.main()