From aa36f83c1670f1e41fa9432a20e5c4a88ee9012c Mon Sep 17 00:00:00 2001
From: Victor Stinner <vstinner@python.org>
Date: Fri, 10 May 2024 21:08:24 +0200
Subject: [PATCH] gh-118702: Implement vectorcall for BaseException (#118703)

* BaseException_vectorcall() now creates a tuple from 'args' array.
* Creation an exception using BaseException_vectorcall() is now a
  single function call, rather than having to call
  BaseException_new() and then BaseException_init().
  Calling BaseException_init() is inefficient since it overrides
  the 'args' attribute.
* _PyErr_SetKeyError() now uses PyObject_CallOneArg() to create the
  KeyError instance to use BaseException_vectorcall().
---
 Lib/test/test_exceptions.py | 23 +++++++++++++++++++++
 Objects/exceptions.c        | 40 +++++++++++++++++++++++++++++++++++++
 Python/errors.c             |  9 +++++----
 3 files changed, 68 insertions(+), 4 deletions(-)

diff --git a/Lib/test/test_exceptions.py b/Lib/test/test_exceptions.py
index 3138f50076f..9460d1f1c86 100644
--- a/Lib/test/test_exceptions.py
+++ b/Lib/test/test_exceptions.py
@@ -1817,6 +1817,29 @@ class ExceptionTests(unittest.TestCase):
         rc, _, err = script_helper.assert_python_ok("-c", code)
         self.assertIn(b'MemoryError', err)
 
+    def test_keyerror_context(self):
+        # Make sure that _PyErr_SetKeyError() chains exceptions
+        try:
+            err1 = None
+            err2 = None
+            try:
+                d = {}
+                try:
+                    raise ValueError("bug")
+                except Exception as exc:
+                    err1 = exc
+                    d[1]
+            except Exception as exc:
+                err2 = exc
+
+            self.assertIsInstance(err1, ValueError)
+            self.assertIsInstance(err2, KeyError)
+            self.assertEqual(err2.__context__, err1)
+        finally:
+            # Break any potential reference cycle
+            exc1 = None
+            exc2 = None
+
 
 class NameErrorTests(unittest.TestCase):
     def test_name_error_has_name(self):
diff --git a/Objects/exceptions.c b/Objects/exceptions.c
index 63c461d34fb..f9cd577c1c1 100644
--- a/Objects/exceptions.c
+++ b/Objects/exceptions.c
@@ -78,6 +78,40 @@ BaseException_init(PyBaseExceptionObject *self, PyObject *args, PyObject *kwds)
     return 0;
 }
 
+
+static PyObject *
+BaseException_vectorcall(PyObject *type_obj, PyObject * const*args,
+                         size_t nargsf, PyObject *kwnames)
+{
+    PyTypeObject *type = _PyType_CAST(type_obj);
+    if (!_PyArg_NoKwnames(type->tp_name, kwnames)) {
+        return NULL;
+    }
+
+    PyBaseExceptionObject *self;
+    self = (PyBaseExceptionObject *)type->tp_alloc(type, 0);
+    if (!self) {
+        return NULL;
+    }
+
+    // The dict is created on the fly in PyObject_GenericSetAttr()
+    self->dict = NULL;
+    self->notes = NULL;
+    self->traceback = NULL;
+    self->cause = NULL;
+    self->context = NULL;
+    self->suppress_context = 0;
+
+    self->args = _PyTuple_FromArray(args, PyVectorcall_NARGS(nargsf));
+    if (!self->args) {
+        Py_DECREF(self);
+        return NULL;
+    }
+
+    return (PyObject *)self;
+}
+
+
 static int
 BaseException_clear(PyBaseExceptionObject *self)
 {
@@ -486,6 +520,7 @@ static PyTypeObject _PyExc_BaseException = {
     (initproc)BaseException_init, /* tp_init */
     0,                          /* tp_alloc */
     BaseException_new,          /* tp_new */
+    .tp_vectorcall = BaseException_vectorcall,
 };
 /* the CPython API expects exceptions to be (PyObject *) - both a hold-over
 from the previous implementation and also allowing Python objects to be used
@@ -3675,6 +3710,11 @@ _PyExc_InitTypes(PyInterpreterState *interp)
         if (_PyStaticType_InitBuiltin(interp, exc) < 0) {
             return -1;
         }
+        if (exc->tp_new == BaseException_new
+            && exc->tp_init == (initproc)BaseException_init)
+        {
+            exc->tp_vectorcall = BaseException_vectorcall;
+        }
     }
     return 0;
 }
diff --git a/Python/errors.c b/Python/errors.c
index 433253b8f9a..ad6b7dbef07 100644
--- a/Python/errors.c
+++ b/Python/errors.c
@@ -257,13 +257,14 @@ void
 _PyErr_SetKeyError(PyObject *arg)
 {
     PyThreadState *tstate = _PyThreadState_GET();
-    PyObject *tup = PyTuple_Pack(1, arg);
-    if (!tup) {
+    PyObject *exc = PyObject_CallOneArg(PyExc_KeyError, arg);
+    if (!exc) {
         /* caller will expect error to be set anyway */
         return;
     }
-    _PyErr_SetObject(tstate, PyExc_KeyError, tup);
-    Py_DECREF(tup);
+
+    _PyErr_SetObject(tstate, (PyObject*)Py_TYPE(exc), exc);
+    Py_DECREF(exc);
 }
 
 void