gh-146480: Override the exception in _PyErr_SetKeyError()

Lib/test/test_capi/test_misc.py

@@ -928,6 +928,23 @@ def test_tp_bases_slot_none(self):
             _testcapi.create_heapctype_with_none_bases_slot
         )
 
+    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),
+            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,))
+
 
 @requires_limited_api
 class TestHeapTypeRelative(unittest.TestCase):

Modules/_testinternalcapi.c

@@ -2838,6 +2838,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_eval_frame_stats", get_eval_frame_stats, METH_NOARGS, NULL},
@@ -2959,6 +2973,7 @@ static PyMethodDef module_functions[] = {
     {"module_get_gc_hooks", module_get_gc_hooks, 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

@@ -246,13 +246,19 @@ 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.
+ *
+ * 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 */