diff --git a/Lib/test/test_capi/test_exceptions.py b/Lib/test/test_capi/test_exceptions.py index ade55338e63b69..6a32201cce5e4a 100644 --- a/Lib/test/test_capi/test_exceptions.py +++ b/Lib/test/test_capi/test_exceptions.py @@ -13,8 +13,9 @@ from .test_misc import decode_stderr -# Skip this test if the _testcapi module isn't available. +# Skip this test if the _testcapi or _testinternalcapi module isn't available. _testcapi = import_helper.import_module('_testcapi') +_testinternalcapi = import_helper.import_module('_testinternalcapi') NULL = None @@ -108,6 +109,26 @@ def __del__(self): b':7: RuntimeWarning: Testing PyErr_WarnEx', ]) + def test__pyerr_setkeyerror(self): + # Test _PyErr_SetKeyError() + _pyerr_setkeyerror = _testinternalcapi._pyerr_setkeyerror + for arg in ( + "key", + # check that a tuple argument is not unpacked + (1, 2, 3), + # PyErr_SetObject(exc_type, exc_value) uses exc_value if it's + # already an exception, but _PyErr_SetKeyError() always creates a + # new KeyError. + KeyError('arg'), + ): + with self.subTest(arg=arg): + with self.assertRaises(KeyError) as cm: + # Test calling _PyErr_SetKeyError() with an exception set + # to check that the function overrides the current + # exception. + _pyerr_setkeyerror(arg) + self.assertEqual(cm.exception.args, (arg,)) + class Test_FatalError(unittest.TestCase): diff --git a/Modules/_testinternalcapi.c b/Modules/_testinternalcapi.c index 27f2d70e832c0f..76697ee9c41a0e 100644 --- a/Modules/_testinternalcapi.c +++ b/Modules/_testinternalcapi.c @@ -2467,6 +2467,20 @@ test_threadstate_set_stack_protection(PyObject *self, PyObject *Py_UNUSED(args)) } +static PyObject * +_pyerr_setkeyerror(PyObject *self, PyObject *arg) +{ + // Test that _PyErr_SetKeyError() overrides the current exception + // if an exception is set + PyErr_NoMemory(); + + _PyErr_SetKeyError(arg); + + assert(PyErr_Occurred()); + return NULL; +} + + static PyMethodDef module_functions[] = { {"get_configs", get_configs, METH_NOARGS}, {"get_recursion_depth", get_recursion_depth, METH_NOARGS}, @@ -2578,6 +2592,7 @@ static PyMethodDef module_functions[] = { {"set_vectorcall_nop", set_vectorcall_nop, METH_O}, {"test_threadstate_set_stack_protection", test_threadstate_set_stack_protection, METH_NOARGS}, + {"_pyerr_setkeyerror", _pyerr_setkeyerror, METH_O}, {NULL, NULL} /* sentinel */ }; diff --git a/Python/errors.c b/Python/errors.c index 3fd97ea36b3368..13633cb20c419c 100644 --- a/Python/errors.c +++ b/Python/errors.c @@ -247,13 +247,23 @@ PyErr_SetObject(PyObject *exception, PyObject *value) _PyErr_SetObject(tstate, exception, value); } -/* Set a key error with the specified argument, wrapping it in a - * tuple automatically so that tuple keys are not unpacked as the - * exception arguments. */ +/* Set a key error with the specified argument. This function should be used to + * raise a KeyError with an argument instead of PyErr_SetObject(PyExc_KeyError, + * arg) which has a special behavior. PyErr_SetObject() unpacks arg if it's a + * tuple, and it uses arg instead of creating a new exception if arg is an + * exception. + * + * If an exception is already set, override the exception. */ void _PyErr_SetKeyError(PyObject *arg) { PyThreadState *tstate = _PyThreadState_GET(); + + // PyObject_CallOneArg() must not be called with an exception set, + // otherwise _Py_CheckFunctionResult() can fail if the function returned + // a result with an excception set. + _PyErr_Clear(tstate); + PyObject *exc = PyObject_CallOneArg(PyExc_KeyError, arg); if (!exc) { /* caller will expect error to be set anyway */