[3.14] gh-146480: Override the exception in _PyErr_SetKeyError() (GH-146486) (#146511)

gh-146480: Override the exception in _PyErr_SetKeyError() (GH-146486)

If _PyErr_SetKeyError() is called with an exception set, it now
replaces the current exception with KeyError (as expected), instead
of setting a SystemError or failing with a fatal error (in debug
mode).
(cherry picked from commit d4153a9)

Co-authored-by: Victor Stinner <vstinner@python.org>

Author: miss-islington

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'<string>: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):
 

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 */
 };
 

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 */