Restore behavior of PyErr_Restore, and add tests.

python · markshannon · Feb 8, 2023 · Feb 6, 2023 · Feb 6, 2023 · Feb 6, 2023
commit 566f6e1ba3675cd077e0fdd2b2ae25aa8a73938c
@@ -1550,5 +1550,44 @@ def func2(x=None):
         self.do_test(func2)
 
 
+class Test_ErrSetAndRestore(unittest.TestCase):
+
+    def test_err_set_raised(self):
+        with self.assertRaises(ValueError):
+            _testcapi.err_set_raised(ValueError())
+        v = ValueError()
+        try:
+            _testcapi.err_set_raised(v)
+        except ValueError as ex:
+            self.assertIs(v, ex)
+
+    def test_err_restore(self):
+        with self.assertRaises(ValueError):
+            _testcapi.err_restore(ValueError)
+        with self.assertRaises(ValueError):
+            _testcapi.err_restore(ValueError, 1)
+        with self.assertRaises(ValueError):
+            _testcapi.err_restore(ValueError, 1, None)
+        with self.assertRaises(ValueError):
+            _testcapi.err_restore(ValueError, ValueError())
+        try:
+            _testcapi.err_restore(KeyError, "hi")
+        except KeyError as k:
+            self.assertEqual("hi", k.args[0])
+        try:
+            1/0
+        except Exception as e:
+            tb = e.__traceback__
+        with self.assertRaises(ValueError):
+            _testcapi.err_restore(ValueError, 1, tb)
+        with self.assertRaises(TypeError):
+            _testcapi.err_restore(ValueError, 1, 0)
+        try:
+            _testcapi.err_restore(ValueError, 1, tb)
+        except ValueError as v:
+            self.assertEqual(1, v.args[0])
+            self.assertIs(tb, v.__traceback__.tb_next)
+
+
 if __name__ == "__main__":
     unittest.main()
diff --git a/Modules/_testcapimodule.c b/Modules/_testcapimodule.c
@@ -2149,7 +2149,7 @@ dict_get_version(PyObject *self, PyObject *args)
         return NULL;
 
     _Py_COMP_DIAG_PUSH
-    _Py_COMP_DIAG_IGNORE_DEPR_DECLS    
+    _Py_COMP_DIAG_IGNORE_DEPR_DECLS
     version = dict->ma_version_tag;
     _Py_COMP_DIAG_POP
 
@@ -3244,6 +3244,41 @@ function_set_kw_defaults(PyObject *self, PyObject *args)
     Py_RETURN_NONE;
 }
 
+static PyObject *
+err_set_raised(PyObject *self, PyObject *exc)
+{
+    Py_INCREF(exc);
+    PyErr_SetRaisedException(exc);
+    assert(PyErr_Occurred());
+    return NULL;
+}
+
+static PyObject *
+err_restore(PyObject *self, PyObject *args) {
+    PyObject *type = NULL, *value = NULL, *traceback = NULL;
+    switch(PyTuple_Size(args)) {
+        case 3:
+            traceback = PyTuple_GetItem(args, 2);
+            Py_INCREF(traceback);
+            /* fall through */
+        case 2:
+            value = PyTuple_GetItem(args, 1);
+            Py_INCREF(value);
+            /* fall through */
+        case 1:
+            type = PyTuple_GetItem(args, 0);
+            Py_INCREF(type);
+            break;
+        default:
+            PyErr_SetString(PyExc_TypeError,
+                        "wrong number of arguments");
+            return NULL;
+    }
+    PyErr_Restore(type, value, traceback);
+    assert(PyErr_Occurred());
+    return NULL;
+}
+
 static PyObject *test_buildvalue_issue38913(PyObject *, PyObject *);
 
 static PyMethodDef TestMethods[] = {
@@ -3392,6 +3427,8 @@ static PyMethodDef TestMethods[] = {
     {"function_set_defaults", function_set_defaults, METH_VARARGS, NULL},
     {"function_get_kw_defaults", function_get_kw_defaults, METH_O, NULL},
     {"function_set_kw_defaults", function_set_kw_defaults, METH_VARARGS, NULL},
+    {"err_set_raised", err_set_raised, METH_O, NULL},
+    {"err_restore", err_restore, METH_VARARGS, NULL},
     {NULL, NULL} /* sentinel */
 };
 

diff --git a/Python/errors.c b/Python/errors.c
@@ -35,35 +35,75 @@ _PyErr_SetRaisedException(PyThreadState *tstate, PyObject *exc)
     Py_XDECREF(old_exc);
 }
 
+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_Restore(PyThreadState *tstate, PyObject *type, PyObject *value,
                PyObject *traceback)
 {
-#ifdef Py_DEBUG
-    /* Check that we are being passed a normalized exception.
-     *
-     * Exceptions are normalized if all NULL,
-     * or if curexc_type = Py_TYPE(curexc_value) and
-     * curexc_traceback = curexc_value->traceback
-     * and both type and traceback are valid */
-    if (value == NULL) {
-        assert(type == NULL);
+    if (type == NULL) {
+        assert(value == NULL);
         assert(traceback == NULL);
+        _PyErr_SetRaisedException(tstate, NULL);
+        return;
     }
-    else {
-        assert(PyExceptionClass_Check(type));
-        assert(type == (PyObject *)Py_TYPE(value));
+    assert(PyExceptionClass_Check(type));
+    if (value != NULL && type == (PyObject *)Py_TYPE(value)) {
+        /* Already normalized */
         assert(((PyBaseExceptionObject *)value)->traceback != Py_None);
-        assert(traceback == ((PyBaseExceptionObject *)value)->traceback ||
-            (traceback == Py_None &&
-            ((PyBaseExceptionObject *)value)->traceback == NULL)
-        );
     }
-#endif
-
+    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_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_XDECREF(type);
-    Py_XDECREF(traceback);
+    Py_DECREF(type);
 }
 
 void
@@ -94,32 +134,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)
 {