Skip to content

gh-111375: Fix handling of exceptions within @contextmanager-decorated functions #111676

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
wants to merge 6 commits into
base: main
Choose a base branch
from
Open
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
13 changes: 13 additions & 0 deletions Lib/contextlib.py
Original file line number Diff line number Diff line change
Expand Up @@ -106,6 +106,7 @@ class _GeneratorContextManagerBase:
"""Shared functionality for @contextmanager and @asynccontextmanager."""

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

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

def __exit__(self, typ, value, traceback):
# don't keep the stored exception alive unnecessarily
exc_context = self.exc_context
self.exc_context = None
if typ is None:
try:
next(self.gen)
Expand All @@ -159,6 +165,13 @@ def __exit__(self, typ, value, traceback):
# tell if we get the same exception back
value = typ()
try:
# If the generator handles the exception thrown into it, the
# exception context reverts to the actual current exception
# context here. In order to make the context manager behave
# like a normal function we set the current exception context
# to what it was during the context manager's __enter__
# (see gh-111676).
sys._set_exception(exc_context)
Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

TBH, I'm not that comfortable adding this function to the sys module, even with an _. It is an API that shouldn't really exist, we're adding it just to enable a hack in contextmanager, which needs to change the semantics of code.

Is there a way that we can roll this functionality into the gen.throw() function instead? Such as -- add a parameter that indicates that the exception stack needs to be fudged in this way?

Copy link
Contributor Author

@pR0Ps pR0Ps Dec 22, 2023

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Yeah, that makes sense. I don't really have strong opinions on how it's implemented, I only went with sys._set_exception because it seemed like it would work and there is precedent for a function that directly manipulates the current exception state.

With that being said, I don't know if doing it entirely within throw() is possible, at least in this case. Because the exception that needs to be set is the one being handled in the __enter__, I don't think there's enough context when throw() is called to correctly manipulate the current exception.

Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I was thinking you pass the exc_context as a new kwwarg to throw which defaults to None.

(I wasn't aware of exc_clear actually. It was added in 2.3, so quite old. Doesn't seem to be used anywhere in the stdlib at the moment.)

Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Ah, that makes sense. Now that I've got sys._set_exception() working and am therefore (reasonably) sure that the None/NULL exception stuff works properly, I'll take a look at adding a kwarg to throw() that essentially just makes it call PyErr_SetHandledException behind the scenes.

self.gen.throw(value)
except StopIteration as exc:
# Suppress StopIteration *unless* it's the same exception that
Expand Down
92 changes: 92 additions & 0 deletions Lib/test/test_contextlib.py
Original file line number Diff line number Diff line change
Expand Up @@ -306,6 +306,98 @@ def woohoo():
with woohoo():
raise StopIteration

def test_contextmanager_handling_exception_resets_exc_info(self):
# Test that sys.exc_info() is correctly unset after handling the error
# when used within a context manager

@contextmanager
def ctx(reraise=False):
try:
self.assertIsNone(sys.exception())
yield
except:
self.assertIsInstance(sys.exception(), ZeroDivisionError)
if reraise:
raise
else:
self.assertIsNone(sys.exception())
self.assertIsNone(sys.exception())

with ctx():
pass

with ctx():
1/0

with self.assertRaises(ZeroDivisionError):
with ctx(reraise=True):
1/0

def test_contextmanager_while_handling(self):
# test that any exceptions currently being handled are preserved
# through the context manager

@contextmanager
def ctx(reraise=False):
# called while handling an IndexError --> TypeError
self.assertIsInstance(sys.exception(), TypeError)
self.assertIsInstance(sys.exception().__context__, IndexError)
exc_ctx = sys.exception()
try:
# raises a ValueError --> ZeroDivisionError
yield
except:
self.assertIsInstance(sys.exception(), ZeroDivisionError)
self.assertIsInstance(sys.exception().__context__, ValueError)
# original error context is preserved
self.assertIs(sys.exception().__context__.__context__, exc_ctx)
if reraise:
raise

# inner error handled, context should now be the original context
self.assertIs(sys.exception(), exc_ctx)

try:
raise IndexError()
except:
try:
raise TypeError()
except:
with ctx():
self.assertIsInstance(sys.exception(), TypeError)
self.assertIsInstance(sys.exception().__context__, IndexError)
try:
raise ValueError()
except:
self.assertIsInstance(sys.exception(), ValueError)
self.assertIsInstance(sys.exception().__context__, TypeError)
self.assertIsInstance(sys.exception().__context__.__context__, IndexError)
1/0
self.assertIsInstance(sys.exception(), TypeError)
self.assertIsInstance(sys.exception().__context__, IndexError)
self.assertIsInstance(sys.exception(), IndexError)

try:
raise IndexError()
except:
try:
raise TypeError()
except:
with self.assertRaises(ZeroDivisionError):
with ctx(reraise=True):
self.assertIsInstance(sys.exception(), TypeError)
self.assertIsInstance(sys.exception().__context__, IndexError)
try:
raise ValueError()
except:
self.assertIsInstance(sys.exception(), ValueError)
self.assertIsInstance(sys.exception().__context__, TypeError)
self.assertIsInstance(sys.exception().__context__.__context__, IndexError)
1/0
self.assertIsInstance(sys.exception(), TypeError)
self.assertIsInstance(sys.exception().__context__, IndexError)
self.assertIsInstance(sys.exception(), IndexError)

def _create_contextmanager_attribs(self):
def attribs(**kw):
def decorate(func):
Expand Down
125 changes: 125 additions & 0 deletions Lib/test/test_sys.py
Original file line number Diff line number Diff line change
Expand Up @@ -142,6 +142,131 @@ def f():
self.assertIsInstance(e, ValueError)
self.assertIs(exc, e)

class SetExceptionTests(unittest.TestCase):

def tearDown(self):
# make sure we don't leave the global exception set
sys._set_exception(None);

def test_set_exc_invalid_values(self):
for x in (0, "1", b"2"):
with self.assertRaises(TypeError):
sys._set_exception(x);

def test_clear_exc(self):
try:
raise ValueError()
except ValueError:
self.assertIsInstance(sys.exception(), ValueError)
sys._set_exception(None)
self.assertIsNone(sys.exception())

def test_set_exc(self):
exc = ValueError()
self.assertIsNone(sys.exception())
sys._set_exception(exc)
self.assertIs(sys.exception(), exc)

def test_set_exc_replaced_by_new_exception_and_restored(self):
exc = ValueError()
sys._set_exception(exc)
self.assertIs(sys.exception(), exc)
try:
raise TypeError()
except TypeError:
self.assertIsInstance(sys.exception(), TypeError)
self.assertIs(sys.exception().__context__, exc)

self.assertIs(sys.exception(), exc)

def test_set_exc_popped_on_exit_except(self):
exc = ValueError()
try:
raise TypeError()
except TypeError:
self.assertIsInstance(sys.exception(), TypeError)
sys._set_exception(exc)
self.assertIs(sys.exception(), exc)
self.assertIsNone(sys.exception())

def test_cleared_exc_overridden_and_restored(self):
try:
raise ValueError()
except ValueError:
try:
raise TypeError()
except TypeError:
self.assertIsInstance(sys.exception(), TypeError)
sys._set_exception(None)
self.assertIsNone(sys.exception())
try:
raise IndexError()
except IndexError:
self.assertIsInstance(sys.exception(), IndexError)
self.assertIsNone(sys.exception().__context__)
self.assertIsNone(sys.exception())
self.assertIsInstance(sys.exception(), ValueError)
self.assertIsNone(sys.exception())

def test_clear_exc_in_generator(self):
def inner():
self.assertIsNone(sys.exception())
yield
self.assertIsInstance(sys.exception(), ValueError)
sys._set_exception(None)
self.assertIsNone(sys.exception())
yield
self.assertIsNone(sys.exception())

# with a single exception in exc_info stack
g = inner()
next(g)
try:
raise ValueError()
except:
self.assertIsInstance(sys.exception(), ValueError)
next(g)
self.assertIsInstance(sys.exception(), ValueError)
self.assertIsNone(sys.exception())
with self.assertRaises(StopIteration):
next(g)
self.assertIsNone(sys.exception())

# with multiple exceptions in exc_info stack by chaining generators
def outer():
g = inner()
self.assertIsNone(sys.exception())
yield next(g)
self.assertIsInstance(sys.exception(), TypeError)
try:
raise ValueError()
except:
self.assertIsInstance(sys.exception(), ValueError)
self.assertIsInstance(sys.exception().__context__, TypeError)
yield next(g)
# at this point the TypeError from the caller has been handled
# by the caller's except block. Even still, it should still be
# referenced as the __context__ of the current exception.
self.assertIsInstance(sys.exception(), ValueError)
self.assertIsInstance(sys.exception().__context__, TypeError)
# not handling an exception, caller isn't handling one either
self.assertIsNone(sys.exception())
with self.assertRaises(StopIteration):
next(g)
self.assertIsNone(sys.exception())

g = outer()
next(g)
try:
raise TypeError()
except:
self.assertIsInstance(sys.exception(), TypeError)
next(g)
self.assertIsInstance(sys.exception(), TypeError)
self.assertIsNone(sys.exception())
with self.assertRaises(StopIteration):
next(g)
self.assertIsNone(sys.exception())

class ExceptHookTest(unittest.TestCase):

Expand Down
Original file line number Diff line number Diff line change
@@ -0,0 +1,2 @@
Add sys._set_exception() function that can set/clear the current exception
context. Patch by Carey Metcalfe.
Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

We have 3 news files here. This and the next one should be deleted.

Copy link
Contributor Author

@pR0Ps pR0Ps Dec 22, 2023

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Will do. Should the sys._set_exception() be added to the other news file or just not noted at all since it's an "internal" function? (provided sys._set_exception() is how this ends up being implemented, otherwise this is a moot point)

Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Probably not needed if it's private. But you can wait with the news file till we know what's going in.

Original file line number Diff line number Diff line change
@@ -0,0 +1,2 @@
Fix handling of ``sys.exception()`` within ``@contextlib.contextmanager``
functions. Patch by Carey Metcalfe.
62 changes: 61 additions & 1 deletion Python/clinic/sysmodule.c.h

Some generated files are not rendered by default. Learn more about how customized files appear on GitHub.

3 changes: 1 addition & 2 deletions Python/errors.c
Original file line number Diff line number Diff line change
Expand Up @@ -125,7 +125,6 @@ _PyErr_GetTopmostException(PyThreadState *tstate)
{
exc_info = exc_info->previous_item;
}
assert(!Py_IsNone(exc_info->exc_value));
return exc_info;
}

Expand Down Expand Up @@ -592,7 +591,7 @@ PyErr_GetHandledException(void)
void
_PyErr_SetHandledException(PyThreadState *tstate, PyObject *exc)
{
Py_XSETREF(tstate->exc_info->exc_value, Py_XNewRef(exc == Py_None ? NULL : exc));
Py_XSETREF(tstate->exc_info->exc_value, Py_XNewRef(exc ? exc : Py_None));
}

void
Expand Down
Loading