diff options
| -rw-r--r-- | Lib/contextlib.py | 5 | ||||
| -rw-r--r-- | Lib/test/test_contextlib.py | 22 |
2 files changed, 26 insertions, 1 deletions
diff --git a/Lib/contextlib.py b/Lib/contextlib.py index 0a5d608..282fc51 100644 --- a/Lib/contextlib.py +++ b/Lib/contextlib.py @@ -30,9 +30,12 @@ class GeneratorContextManager(object): else: try: self.gen.throw(type, value, traceback) - return True + raise RuntimeError("generator didn't stop after throw()") except StopIteration: return True + except: + if sys.exc_info()[1] is not value: + raise def contextmanager(func): diff --git a/Lib/test/test_contextlib.py b/Lib/test/test_contextlib.py index f8db88c..7d7f8d2 100644 --- a/Lib/test/test_contextlib.py +++ b/Lib/test/test_contextlib.py @@ -45,6 +45,28 @@ class ContextManagerTestCase(unittest.TestCase): self.fail("Expected ZeroDivisionError") self.assertEqual(state, [1, 42, 999]) + def test_contextmanager_no_reraise(self): + @contextmanager + def whee(): + yield + ctx = whee().__context__() + ctx.__enter__() + # Calling __exit__ should not result in an exception + self.failIf(ctx.__exit__(TypeError, TypeError("foo"), None)) + + def test_contextmanager_trap_yield_after_throw(self): + @contextmanager + def whoo(): + try: + yield + except: + yield + ctx = whoo().__context__() + ctx.__enter__() + self.assertRaises( + RuntimeError, ctx.__exit__, TypeError, TypeError("foo"), None + ) + def test_contextmanager_except(self): state = [] @contextmanager |