10 files changed, 171 insertions, 120 deletions

diff --git a/Python/bytecodes.c b/Python/bytecodes.c
index 9633f34212a..1169d8d172d 100644
--- a/Python/bytecodes.c
+++ b/Python/bytecodes.c
@@ -804,9 +804,7 @@ dummy_func(
                 DECREF_INPUTS();
             }
             else {
-                PyObject *exc_type = Py_NewRef(Py_TYPE(exc_value));
-                PyObject *exc_traceback = PyException_GetTraceback(exc_value);
-                _PyErr_Restore(tstate, exc_type, Py_NewRef(exc_value), exc_traceback);
+                _PyErr_SetRaisedException(tstate, Py_NewRef(exc_value));
                 goto exception_unwind;
             }
         }
diff --git a/Python/ceval.c b/Python/ceval.c
index ecb5bf96555..a91f5baca88 100644
--- a/Python/ceval.c
+++ b/Python/ceval.c
@@ -2902,13 +2902,13 @@ format_kwargs_error(PyThreadState *tstate, PyObject *func, PyObject *kwargs)
         }
     }
     else if (_PyErr_ExceptionMatches(tstate, PyExc_KeyError)) {
-        PyObject *exc, *val, *tb;
-        _PyErr_Fetch(tstate, &exc, &val, &tb);
-        if (val && PyTuple_Check(val) && PyTuple_GET_SIZE(val) == 1) {
+        PyObject *exc = _PyErr_GetRaisedException(tstate);
+        PyObject *args = ((PyBaseExceptionObject *)exc)->args;
+        if (exc && PyTuple_Check(args) && PyTuple_GET_SIZE(args) == 1) {
             _PyErr_Clear(tstate);
             PyObject *funcstr = _PyObject_FunctionStr(func);
             if (funcstr != NULL) {
-                PyObject *key = PyTuple_GET_ITEM(val, 0);
+                PyObject *key = PyTuple_GET_ITEM(args, 0);
                 _PyErr_Format(
                     tstate, PyExc_TypeError,
                     "%U got multiple values for keyword argument '%S'",
@@ -2916,11 +2916,9 @@ format_kwargs_error(PyThreadState *tstate, PyObject *func, PyObject *kwargs)
                 Py_DECREF(funcstr);
             }
             Py_XDECREF(exc);
-            Py_XDECREF(val);
-            Py_XDECREF(tb);
         }
         else {
-            _PyErr_Restore(tstate, exc, val, tb);
+            _PyErr_SetRaisedException(tstate, exc);
         }
     }
 }
diff --git a/Python/errors.c b/Python/errors.c
index 05ef62246ec..f573bed3d63 100644
--- a/Python/errors.c
+++ b/Python/errors.c
@@ -27,32 +27,84 @@ static PyObject *
 _PyErr_FormatV(PyThreadState *tstate, PyObject *exception,
                const char *format, va_list vargs);
 
-
 void
-_PyErr_Restore(PyThreadState *tstate, PyObject *type, PyObject *value,
-               PyObject *traceback)
+_PyErr_SetRaisedException(PyThreadState *tstate, PyObject *exc)
 {
-    PyObject *oldtype, *oldvalue, *oldtraceback;
+    PyObject *old_exc = tstate->current_exception;
+    tstate->current_exception = exc;
+    Py_XDECREF(old_exc);
+}
 
-    if (traceback != NULL && !PyTraceBack_Check(traceback)) {
-        /* XXX Should never happen -- fatal error instead? */
-        /* Well, it could be None. */
-        Py_SETREF(traceback, NULL);
+static PyObject*
+_PyErr_CreateException(PyObject *exception_type, PyObject *value)
+{
+    PyObject *exc;
+
+    if (value == NULL || value == Py_None) {
+        exc = _PyObject_CallNoArgs(exception_type);
+    }
+    else if (PyTuple_Check(value)) {
+        exc = PyObject_Call(exception_type, value, NULL);
+    }
+    else {
+        exc = PyObject_CallOneArg(exception_type, value);
     }
 
-    /* Save these in locals to safeguard against recursive
-       invocation through Py_XDECREF */
-    oldtype = tstate->curexc_type;
-    oldvalue = tstate->curexc_value;
-    oldtraceback = tstate->curexc_traceback;
+    if (exc != NULL && !PyExceptionInstance_Check(exc)) {
+        PyErr_Format(PyExc_TypeError,
+                     "calling %R should have returned an instance of "
+                     "BaseException, not %s",
+                     exception_type, Py_TYPE(exc)->tp_name);
+        Py_CLEAR(exc);
+    }
 
-    tstate->curexc_type = type;
-    tstate->curexc_value = value;
-    tstate->curexc_traceback = traceback;
+    return exc;
+}
 
-    Py_XDECREF(oldtype);
-    Py_XDECREF(oldvalue);
-    Py_XDECREF(oldtraceback);
+void
+_PyErr_Restore(PyThreadState *tstate, PyObject *type, PyObject *value,
+               PyObject *traceback)
+{
+    if (type == NULL) {
+        assert(value == NULL);
+        assert(traceback == NULL);
+        _PyErr_SetRaisedException(tstate, NULL);
+        return;
+    }
+    assert(PyExceptionClass_Check(type));
+    if (value != NULL && type == (PyObject *)Py_TYPE(value)) {
+        /* Already normalized */
+        assert(((PyBaseExceptionObject *)value)->traceback != Py_None);
+    }
+    else {
+        PyObject *exc = _PyErr_CreateException(type, value);
+        Py_XDECREF(value);
+        if (exc == NULL) {
+            Py_DECREF(type);
+            Py_XDECREF(traceback);
+            return;
+        }
+        value = exc;
+    }
+    assert(PyExceptionInstance_Check(value));
+    if (traceback != NULL && !PyTraceBack_Check(traceback)) {
+        if (traceback == Py_None) {
+            Py_DECREF(Py_None);
+            traceback = NULL;
+        }
+        else {
+            PyErr_SetString(PyExc_TypeError, "traceback must be a Traceback or None");
+            Py_XDECREF(value);
+            Py_DECREF(type);
+            Py_XDECREF(traceback);
+            return;
+        }
+    }
+    PyObject *old_traceback = ((PyBaseExceptionObject *)value)->traceback;
+    ((PyBaseExceptionObject *)value)->traceback = traceback;
+    Py_XDECREF(old_traceback);
+    _PyErr_SetRaisedException(tstate, value);
+    Py_DECREF(type);
 }
 
 void
@@ -62,6 +114,12 @@ PyErr_Restore(PyObject *type, PyObject *value, PyObject *traceback)
     _PyErr_Restore(tstate, type, value, traceback);
 }
 
+void
+PyErr_SetRaisedException(PyObject *exc)
+{
+    PyThreadState *tstate = _PyThreadState_GET();
+    _PyErr_SetRaisedException(tstate, exc);
+}
 
 _PyErr_StackItem *
 _PyErr_GetTopmostException(PyThreadState *tstate)
@@ -77,32 +135,6 @@ _PyErr_GetTopmostException(PyThreadState *tstate)
     return exc_info;
 }
 
-static PyObject*
-_PyErr_CreateException(PyObject *exception_type, PyObject *value)
-{
-    PyObject *exc;
-
-    if (value == NULL || value == Py_None) {
-        exc = _PyObject_CallNoArgs(exception_type);
-    }
-    else if (PyTuple_Check(value)) {
-        exc = PyObject_Call(exception_type, value, NULL);
-    }
-    else {
-        exc = PyObject_CallOneArg(exception_type, value);
-    }
-
-    if (exc != NULL && !PyExceptionInstance_Check(exc)) {
-        PyErr_Format(PyExc_TypeError,
-                     "calling %R should have returned an instance of "
-                     "BaseException, not %s",
-                     exception_type, Py_TYPE(exc)->tp_name);
-        Py_CLEAR(exc);
-    }
-
-    return exc;
-}
-
 void
 _PyErr_SetObject(PyThreadState *tstate, PyObject *exception, PyObject *value)
 {
@@ -117,30 +149,29 @@ _PyErr_SetObject(PyThreadState *tstate, PyObject *exception, PyObject *value)
                       exception);
         return;
     }
-
     Py_XINCREF(value);
-    exc_value = _PyErr_GetTopmostException(tstate)->exc_value;
-    if (exc_value != NULL && exc_value != Py_None) {
-        /* Implicit exception chaining */
-        Py_INCREF(exc_value);
-        if (value == NULL || !PyExceptionInstance_Check(value)) {
-            /* We must normalize the value right now */
-            PyObject *fixed_value;
-
-            /* Issue #23571: functions must not be called with an
-               exception set */
-            _PyErr_Clear(tstate);
+    /* Normalize the exception */
+    if (value == NULL || (PyObject *)Py_TYPE(value) != exception) {
+        /* We must normalize the value right now */
+        PyObject *fixed_value;
 
-            fixed_value = _PyErr_CreateException(exception, value);
-            Py_XDECREF(value);
-            if (fixed_value == NULL) {
-                Py_DECREF(exc_value);
-                return;
-            }
+        /* Issue #23571: functions must not be called with an
+            exception set */
+        _PyErr_Clear(tstate);
 
-            value = fixed_value;
+        fixed_value = _PyErr_CreateException(exception, value);
+        Py_XDECREF(value);
+        if (fixed_value == NULL) {
+            return;
         }
 
+        value = fixed_value;
+    }
+
+    exc_value = _PyErr_GetTopmostException(tstate)->exc_value;
+    if (exc_value != NULL && exc_value != Py_None) {
+        /* Implicit exception chaining */
+        Py_INCREF(exc_value);
         /* Avoid creating new reference cycles through the
            context chain, while taking care not to hang on
            pre-existing ones.
@@ -414,17 +445,34 @@ PyErr_NormalizeException(PyObject **exc, PyObject **val, PyObject **tb)
 }
 
 
+PyObject *
+_PyErr_GetRaisedException(PyThreadState *tstate) {
+    PyObject *exc = tstate->current_exception;
+    tstate->current_exception = NULL;
+    return exc;
+}
+
+PyObject *
+PyErr_GetRaisedException(void)
+{
+    PyThreadState *tstate = _PyThreadState_GET();
+    return _PyErr_GetRaisedException(tstate);
+}
+
 void
 _PyErr_Fetch(PyThreadState *tstate, PyObject **p_type, PyObject **p_value,
              PyObject **p_traceback)
 {
-    *p_type = tstate->curexc_type;
-    *p_value = tstate->curexc_value;
-    *p_traceback = tstate->curexc_traceback;
-
-    tstate->curexc_type = NULL;
-    tstate->curexc_value = NULL;
-    tstate->curexc_traceback = NULL;
+    PyObject *exc = _PyErr_GetRaisedException(tstate);
+    *p_value = exc;
+    if (exc == NULL) {
+        *p_type = NULL;
+        *p_traceback = NULL;
+    }
+    else {
+        *p_type = Py_NewRef(Py_TYPE(exc));
+        *p_traceback = Py_XNewRef(((PyBaseExceptionObject *)exc)->traceback);
+    }
 }
 
 
@@ -597,6 +645,28 @@ _PyErr_ChainExceptions(PyObject *typ, PyObject *val, PyObject *tb)
     }
 }
 
+/* Like PyErr_SetRaisedException(), but if an exception is already set,
+   set the context associated with it.
+
+   The caller is responsible for ensuring that this call won't create
+   any cycles in the exception context chain. */
+void
+_PyErr_ChainExceptions1(PyObject *exc)
+{
+    if (exc == NULL) {
+        return;
+    }
+    PyThreadState *tstate = _PyThreadState_GET();
+    if (_PyErr_Occurred(tstate)) {
+        PyObject *exc2 = _PyErr_GetRaisedException(tstate);
+        PyException_SetContext(exc2, exc);
+        _PyErr_SetRaisedException(tstate, exc2);
+    }
+    else {
+        _PyErr_SetRaisedException(tstate, exc);
+    }
+}
+
 /* Set the currently set exception's context to the given exception.
 
    If the provided exc_info is NULL, then the current Python thread state's
@@ -707,19 +777,6 @@ PyErr_BadArgument(void)
 }
 
 PyObject *
-_PyErr_NoMemory(PyThreadState *tstate)
-{
-    if (Py_IS_TYPE(PyExc_MemoryError, NULL)) {
-        /* PyErr_NoMemory() has been called before PyExc_MemoryError has been
-           initialized by _PyExc_Init() */
-        Py_FatalError("Out of memory and PyExc_MemoryError is not "
-                      "initialized yet");
-    }
-    _PyErr_SetNone(tstate, PyExc_MemoryError);
-    return NULL;
-}
-
-PyObject *
 PyErr_NoMemory(void)
 {
     PyThreadState *tstate = _PyThreadState_GET();
diff --git a/Python/generated_cases.c.h b/Python/generated_cases.c.h
index f38286441be..09eb6893ebf 100644
--- a/Python/generated_cases.c.h
+++ b/Python/generated_cases.c.h
@@ -1036,9 +1036,7 @@
                 Py_DECREF(exc_value);
             }
             else {
-                PyObject *exc_type = Py_NewRef(Py_TYPE(exc_value));
-                PyObject *exc_traceback = PyException_GetTraceback(exc_value);
-                _PyErr_Restore(tstate, exc_type, Py_NewRef(exc_value), exc_traceback);
+                _PyErr_SetRaisedException(tstate, Py_NewRef(exc_value));
                 goto exception_unwind;
             }
             STACK_SHRINK(2);
diff --git a/Python/import.c b/Python/import.c
index da6c15c5fd4..1318c09d9b3 100644
--- a/Python/import.c
+++ b/Python/import.c
@@ -1592,6 +1592,13 @@ remove_importlib_frames(PyThreadState *tstate)
         Py_DECREF(code);
         tb = next;
     }
+    assert(PyExceptionInstance_Check(value));
+    assert((PyObject *)Py_TYPE(value) == exception);
+    if (base_tb == NULL) {
+        base_tb = Py_None;
+        Py_INCREF(Py_None);
+    }
+    PyException_SetTraceback(value, base_tb);
 done:
     _PyErr_Restore(tstate, exception, value, base_tb);
 }
diff --git a/Python/initconfig.c b/Python/initconfig.c
index d7b2dc4a297..deec805a6b1 100644
--- a/Python/initconfig.c
+++ b/Python/initconfig.c
@@ -3143,8 +3143,7 @@ init_dump_ascii_wstr(const wchar_t *str)
 void
 _Py_DumpPathConfig(PyThreadState *tstate)
 {
-    PyObject *exc_type, *exc_value, *exc_tb;
-    _PyErr_Fetch(tstate, &exc_type, &exc_value, &exc_tb);
+    PyObject *exc = _PyErr_GetRaisedException(tstate);
 
     PySys_WriteStderr("Python path configuration:\n");
 
@@ -3202,5 +3201,5 @@ _Py_DumpPathConfig(PyThreadState *tstate)
         PySys_WriteStderr("  ]\n");
     }
 
-    _PyErr_Restore(tstate, exc_type, exc_value, exc_tb);
+    _PyErr_SetRaisedException(tstate, exc);
 }
diff --git a/Python/pystate.c b/Python/pystate.c
index ed8c2e212a5..1261092d143 100644
--- a/Python/pystate.c
+++ b/Python/pystate.c
@@ -1375,9 +1375,7 @@ PyThreadState_Clear(PyThreadState *tstate)
     Py_CLEAR(tstate->dict);
     Py_CLEAR(tstate->async_exc);
 
-    Py_CLEAR(tstate->curexc_type);
-    Py_CLEAR(tstate->curexc_value);
-    Py_CLEAR(tstate->curexc_traceback);
+    Py_CLEAR(tstate->current_exception);
 
     Py_CLEAR(tstate->exc_state.exc_value);
 
diff --git a/Python/pythonrun.c b/Python/pythonrun.c
index 35292b6478a..6a4d5937686 100644
--- a/Python/pythonrun.c
+++ b/Python/pythonrun.c
@@ -748,13 +748,10 @@ _Py_HandleSystemExit(int *exitcode_p)
     }
 
  done:
-    /* Restore and clear the exception info, in order to properly decref
-     * the exception, value, and traceback.      If we just exit instead,
-     * these leak, which confuses PYTHONDUMPREFS output, and may prevent
-     * some finalizers from running.
-     */
-    PyErr_Restore(exception, value, tb);
-    PyErr_Clear();
+    /* Cleanup the exception */
+    Py_CLEAR(exception);
+    Py_CLEAR(value);
+    Py_CLEAR(tb);
     *exitcode_p = exitcode;
     return 1;
 }
diff --git a/Python/sysmodule.c b/Python/sysmodule.c
index f9f766a94d1..6e81ef92b67 100644
--- a/Python/sysmodule.c
+++ b/Python/sysmodule.c
@@ -66,12 +66,11 @@ _PySys_GetAttr(PyThreadState *tstate, PyObject *name)
     if (sd == NULL) {
         return NULL;
     }
-    PyObject *exc_type, *exc_value, *exc_tb;
-    _PyErr_Fetch(tstate, &exc_type, &exc_value, &exc_tb);
+    PyObject *exc = _PyErr_GetRaisedException(tstate);
     /* XXX Suppress a new exception if it was raised and restore
      * the old one. */
     PyObject *value = _PyDict_GetItemWithError(sd, name);
-    _PyErr_Restore(tstate, exc_type, exc_value, exc_tb);
+    _PyErr_SetRaisedException(tstate, exc);
     return value;
 }
 
@@ -3704,11 +3703,10 @@ static void
 sys_format(PyObject *key, FILE *fp, const char *format, va_list va)
 {
     PyObject *file, *message;
-    PyObject *error_type, *error_value, *error_traceback;
     const char *utf8;
     PyThreadState *tstate = _PyThreadState_GET();
 
-    _PyErr_Fetch(tstate, &error_type, &error_value, &error_traceback);
+    PyObject *error = _PyErr_GetRaisedException(tstate);
     file = _PySys_GetAttr(tstate, key);
     message = PyUnicode_FromFormatV(format, va);
     if (message != NULL) {
@@ -3720,7 +3718,7 @@ sys_format(PyObject *key, FILE *fp, const char *format, va_list va)
         }
         Py_DECREF(message);
     }
-    _PyErr_Restore(tstate, error_type, error_value, error_traceback);
+    _PyErr_SetRaisedException(tstate, error);
 }
 
 void
diff --git a/Python/traceback.c b/Python/traceback.c
index da26c9b260a..31b85e77575 100644
--- a/Python/traceback.c
+++ b/Python/traceback.c
@@ -249,6 +249,8 @@ PyTraceBack_Here(PyFrameObject *frame)
         _PyErr_ChainExceptions(exc, val, tb);
         return -1;
     }
+    assert(PyExceptionInstance_Check(val));
+    PyException_SetTraceback(val, newtb);
     PyErr_Restore(exc, val, newtb);
     Py_XDECREF(tb);
     return 0;
@@ -260,13 +262,12 @@ void _PyTraceback_Add(const char *funcname, const char *filename, int lineno)
     PyObject *globals;
     PyCodeObject *code;
     PyFrameObject *frame;
-    PyObject *exc, *val, *tb;
     PyThreadState *tstate = _PyThreadState_GET();
 
     /* Save and clear the current exception. Python functions must not be
        called with an exception set. Calling Python functions happens when
        the codec of the filesystem encoding is implemented in pure Python. */
-    _PyErr_Fetch(tstate, &exc, &val, &tb);
+    PyObject *exc = _PyErr_GetRaisedException(tstate);
 
     globals = PyDict_New();
     if (!globals)
@@ -283,13 +284,13 @@ void _PyTraceback_Add(const char *funcname, const char *filename, int lineno)
         goto error;
     frame->f_lineno = lineno;
 
-    _PyErr_Restore(tstate, exc, val, tb);
+    _PyErr_SetRaisedException(tstate, exc);
     PyTraceBack_Here(frame);
     Py_DECREF(frame);
     return;
 
 error:
-    _PyErr_ChainExceptions(exc, val, tb);
+    _PyErr_ChainExceptions1(exc);
 }
 
 static PyObject *