8000 Fix exception context in @contextmanagers · python/cpython@a6f948b · GitHub
[go: up one dir, main page]
Skip to content

Commit a6f948b

Browse files
pR0PspR0Ps
committed
Fix exception context in @contextmanagers
Previously, when using a `try...yield...except` construct within a `@contextmanager` function, an exception raised by the `yield` wouldn't be properly cleared when handled the `except` block. Instead, calling `sys.exception()` would continue to return the exception, leading to confusing tracebacks and logs. This was happening due to the way that the `@contextmanager` decorator drives its decorated generator function. When an exception occurs, it uses `.throw(exc)` to throw the exception into the generator. However, even if the generator handles this exception, because the exception was thrown into it in the context of the `@contextmanager` decorator handling it (and not being finished yet), `sys.exception()` was not being reset. In order to fix this, the exception context as the `@contextmanager` decorator is `__enter__`'d is stored and set as the current exception context just before throwing the new exception into the generator. Doing this means that after the generator handles the thrown exception, `sys.exception()` reverts back to what it was when the `@contextmanager` function was started.
1 parent d2f6251 commit a6f948b

File tree

3 files changed

+96
-0
lines changed

3 files changed

+96
-0
lines changed

‎Lib/contextlib.py

Lines changed: 12 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -106,6 +106,7 @@ class _GeneratorContextManagerBase:
106106
"""Shared functionality for @contextmanager and @asynccontextmanager."""
107107

108108
def __init__(self, func, args, kwds):
109+
self.exc_context = None
109110
self.gen = func(*args, **kwds)
110111
self.func, self.args, self.kwds = func, args, kwds
111112
# Issue 19330: ensure context manager instances have good docstrings
@@ -134,6 +135,8 @@ class _GeneratorContextManager(
134135
"""Helper for @contextmanager decorator."""
135136

136137
def __enter__(self):
138+
# store the exception context on enter so it can be restored on exit
139+
self.exc_context = sys.exception()
137140
# do not keep args and kwds alive unnecessarily
138141
# they are only n 8000 eeded for recreation, which is not possible anymore
139142
del self.args, self.kwds, self.func
@@ -143,6 +146,9 @@ def __enter__(self):
143146
raise RuntimeError("generator didn't yield") from None
144147

145148
def __exit__(self, typ, value, traceback):
149+
# don't keep the stored exception alive unnecessarily
150+
exc_context = self.exc_context
151+
self.exc_context = None
146152
if typ is None:
147153
try:
148154
next(self.gen)
@@ -159,6 +165,12 @@ def __exit__(self, typ, value, traceback):
159165
# tell if we get the same exception back
160166
value = typ()
161167
try:
168+
# If the generator handles the exception thrown into it, the
169+
# exception context will revert to the actual current exception
170+
# context here. In order to make the context manager behave
171+
# like a normal function we set the current exception context
172+
# to what it was during the context manager's __enter__
173+
sys._set_exception(exc_context)
162174
self.gen.throw(value)
163175
except StopIteration as exc:
164176
# Suppress StopIteration *unless* it's the same exception that

‎Lib/test/test_contextlib.py

Lines changed: 82 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -306,6 +306,88 @@ def woohoo():
306306
with woohoo():
307307
raise StopIteration
308308

309+
def test_contextmanager_handling_exception_resets_exc_info(self):
310+
# Test that sys.exc_info() is correctly unset after handling the error
311+
# when used within a context manager
312+
313+
@contextmanager
314+
def ctx(reraise=False):
315+
try:
316+
self.assertIsNone(sys.exception())
317+
yield
318+
except:
319+
self.assertIsInstance(sys.exception(), ZeroDivisionError)
320+
if reraise:
321+
raise
322+
else:
323+
self.assertIsNone(sys.exception())
324+
self.assertIsNone(sys.exception())
325+
326+
with ctx():
327+
pass
328+
329+
with ctx():
330+
1/0
331+
332+
with self.assertRaises(ZeroDivisionError):
333+
with ctx(reraise=True):
334+
1/0
335+
336+
def test_contextmanager_while_handling(self):
337+
# test that any exceptions currently being handled are preserved
338+
# through the context manager
339+
340+
@contextmanager
341+
def ctx(reraise=False):
342+
# called while handling an IndexError --> TypeError
343+
self.assertIsInstance(sys.exception(), TypeError)
344+
self.assertIsInstance(sys.exception().__context__, IndexError)
345+
exc_ctx = sys.exception()
346+
try:
347+
# raises a ValueError --> ZeroDivisionError
348+
yield
349+
except:
350+
self.assertIsInstance(sys.exception(), ZeroDivisionError)
351+
self.assertIsInstance(sys.exception().__context__, ValueError)
352+
# original error context is preserved
353+
self.assertIs(sys.exception().__context__.__context__, exc_ctx)
354+
if reraise:
355+
raise
356+
357+
# inner error handled, context should now be the original context
358+
self.assertIs(sys.exception(), exc_ctx)
359+
360+
try:
361+
raise IndexError()
362+
except:
363+
try:
364+
raise TypeError()
365+
except:
366+
with ctx():
367+
try:
368+
raise ValueError()
369+
except:
370+
self.assertIsInstance(sys.exception(), ValueError)
371+
1/0
372+
self.assertIsInstance(sys.exception(), TypeError)
373+
self.assertIsInstance(sys.exception(), IndexError)
374+
375+
try:
376+
raise IndexError()
377+
except:
378+
try:
379+
raise TypeError()
380+
except:
381+
with self.assertRaises(ZeroDivisionError):
382+
with ctx(reraise=True):
383+
try:
384+
raise ValueError()
385+
except:
386+
self.assertIsInstance(sys.exception(), ValueError)
387+
1/0
388+
self.assertIsInstance(sys.exception(), TypeError)
389+
self.assertIsInstance(sys.exception(), IndexError)
390+
309391
def _create_contextmanager_attribs(self):
310392
def attribs(**kw):
311393
def decorate(func):

‎Misc/NEWS.d/next/Library/2023-11-03-01-44-51.gh-issue-111375.qDBcCI.rst

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,2 @@
1+
Fix handling of ``sys.exception()`` within ``@contextlib.contextmanager``
2+
functions. Patch by Carey Metcalfe.

0 commit comments

Comments
 (0)
0