gh-146480: Override the exception in _PyErr_SetKeyError()

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

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