From d929f2761e55687da8a222ee13f620613bd8bd3f Mon Sep 17 00:00:00 2001
From: Neil Schemenauer <nas@arctrix.com>
Date: Wed, 31 May 2023 14:21:24 -0700
Subject: [PATCH] wip: add defer_close logic to FileIO

If the file is wrapped by a write buffer, set defer_close to true.  That
ensures the underlying FileIO object is not closed until the buffer is
closed (and can flush the unwritten data).  This avoids data loss, as
demoed by buffer_not_flushed.py.
---
 Modules/_io/bufferedio.c |  9 +++++++++
 Modules/_io/fileio.c     | 13 +++++++++++++
 2 files changed, 22 insertions(+)

diff --git a/Modules/_io/bufferedio.c b/Modules/_io/bufferedio.c
index 44c3f56caee18b6..a805cf5ce602592 100644
--- a/Modules/_io/bufferedio.c
+++ b/Modules/_io/bufferedio.c
@@ -545,6 +545,8 @@ _io__Buffered_closed_get_impl(buffered *self)
     return PyObject_GetAttr(self->raw, &_Py_ID(closed));
 }
 
+void _PyFileIO_set_defer_close(PyObject *self, bool defer_close);
+
 /*[clinic input]
 @critical_section
 _io._Buffered.close
@@ -588,6 +590,10 @@ _io__Buffered_close_impl(buffered *self)
         exc = PyErr_GetRaisedException();
     }
 
+    if (self->fast_closed_checks) {
+        _PyFileIO_set_defer_close(self->raw, false);
+    }
+
     res = PyObject_CallMethodNoArgs(self->raw, &_Py_ID(close));
 
     if (self->buffer) {
@@ -1972,6 +1978,9 @@ _io_BufferedWriter___init___impl(buffered *self, PyObject *raw,
         Py_IS_TYPE(self, state->PyBufferedWriter_Type) &&
         Py_IS_TYPE(raw, state->PyFileIO_Type)
     );
+    if (self->fast_closed_checks) {
+        _PyFileIO_set_defer_close(self->raw, true);
+    }
 
     self->ok = 1;
     return 0;
diff --git a/Modules/_io/fileio.c b/Modules/_io/fileio.c
index e5cd7aca5de2f17..009bbc5acc4ced7 100644
--- a/Modules/_io/fileio.c
+++ b/Modules/_io/fileio.c
@@ -82,6 +82,8 @@ typedef struct {
     struct _Py_stat_struct *stat_atopen;
     PyObject *weakreflist;
     PyObject *dict;
+    /* if true, we are wrapped by buffer, don't close until buffer is closed */
+    bool defer_close;
 } fileio;
 
 #define PyFileIO_Check(state, op) (PyObject_TypeCheck((op), state->PyFileIO_Type))
@@ -96,6 +98,13 @@ _PyFileIO_closed(PyObject *self)
     return (PyFileIO_CAST(self)->fd < 0);
 }
 
+void
+_PyFileIO_set_defer_close(PyObject *self, bool defer_close)
+{
+    PyFileIO_CAST(self)->defer_close = defer_close;
+}
+
+
 /* Because this can call arbitrary code, it shouldn't be called when
    the refcount is 0 (that is, not directly from tp_dealloc unless
    the refcount has been temporarily re-incremented). */
@@ -128,6 +137,9 @@ _io_FileIO__dealloc_warn_impl(fileio *self, PyObject *source)
 static int
 internal_close(fileio *self)
 {
+    if (self->defer_close) {
+        return 0;
+    }
     int err = 0;
     int save_errno = 0;
     if (self->fd >= 0) {
@@ -221,6 +233,7 @@ fileio_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
     self->stat_atopen = NULL;
     self->closefd = 1;
     self->weakreflist = NULL;
+    self->defer_close = false;
     return (PyObject *) self;
 }