flush_io too

iritkatriel · iritkatriel · commit 86af53f5a701 · 2023-03-15T23:31:11.000Z
diff --git a/Python/pythonrun.c b/Python/pythonrun.c
@@ -1639,35 +1639,29 @@ PyRun_FileExFlags(FILE *fp, const char *filename, int start, PyObject *globals,
 
 }
 
-
 static void
-flush_io(void)
+flush_io_stream(PyThreadState *tstate, PyObject *name)
 {
-    PyObject *f, *r;
-    PyObject *type, *value, *traceback;
-
-    /* Save the current exception */
-    PyErr_Fetch(&type, &value, &traceback);
-
-    PyThreadState *tstate = _PyThreadState_GET();
-    f = _PySys_GetAttr(tstate, &_Py_ID(stderr));
-    if (f != NULL) {
-        r = _PyObject_CallMethodNoArgs(f, &_Py_ID(flush));
-        if (r)
-            Py_DECREF(r);
-        else
-            PyErr_Clear();
-    }
-    f = _PySys_GetAttr(tstate, &_Py_ID(stdout));
+    PyObject *f = _PySys_GetAttr(tstate, &_Py_ID(stderr));
     if (f != NULL) {
-        r = _PyObject_CallMethodNoArgs(f, &_Py_ID(flush));
-        if (r)
+        PyObject *r = _PyObject_CallMethodNoArgs(f, &_Py_ID(flush));
+        if (r) {
             Py_DECREF(r);
-        else
+        }
+        else {
             PyErr_Clear();
+        }
     }
+}
 
-    PyErr_Restore(type, value, traceback);
+static void
+flush_io(void)
+{
+    PyThreadState *tstate = _PyThreadState_GET();
+    PyObject *exc = _PyErr_GetRaisedException(tstate);
+    flush_io_stream(tstate, &_Py_ID(stderr));
+    flush_io_stream(tstate, &_Py_ID(stdout));
+    _PyErr_SetRaisedException(tstate, exc);
 }
 
 static PyObject *
