Issue #22982: Improve BOM handling when seeking to multiple positions of a writable text file.

1 files changed, 21 insertions, 4 deletions

diff --git a/Modules/_io/textio.c b/Modules/_io/textio.c
index d1c0d01..b419275 100644
--- a/Modules/_io/textio.c
+++ b/Modules/_io/textio.c
@@ -2042,11 +2042,10 @@ _textiowrapper_decoder_setstate(textio *self, cookie_type *cookie)
 }
 
 static int
-_textiowrapper_encoder_setstate(textio *self, cookie_type *cookie)
+_textiowrapper_encoder_reset(textio *self, int start_of_stream)
 {
     PyObject *res;
-    /* Same as _textiowrapper_decoder_setstate() above. */
-    if (cookie->start_pos == 0 && cookie->dec_flags == 0) {
+    if (start_of_stream) {
         res = PyObject_CallMethodObjArgs(self->encoder, _PyIO_str_reset, NULL);
         self->encoding_start_of_stream = 1;
     }
@@ -2061,6 +2060,14 @@ _textiowrapper_encoder_setstate(textio *self, cookie_type *cookie)
     return 0;
 }
 
+static int
+_textiowrapper_encoder_setstate(textio *self, cookie_type *cookie)
+{
+    /* Same as _textiowrapper_decoder_setstate() above. */
+    return _textiowrapper_encoder_reset(
+        self, cookie->start_pos == 0 && cookie->dec_flags == 0);
+}
+
 static PyObject *
 textiowrapper_seek(textio *self, PyObject *args)
 {
@@ -2128,7 +2135,17 @@ textiowrapper_seek(textio *self, PyObject *args)
         }
 
         res = _PyObject_CallMethodId(self->buffer, &PyId_seek, "ii", 0, 2);
-        Py_XDECREF(cookieObj);
+        Py_CLEAR(cookieObj);
+        if (res == NULL)
+            goto fail;
+        if (self->encoder) {
+            /* If seek() == 0, we are at the start of stream, otherwise not */
+            cmp = PyObject_RichCompareBool(res, _PyIO_zero, Py_EQ);
+            if (cmp < 0 || _textiowrapper_encoder_reset(self, cmp)) {
+                Py_DECREF(res);
+                goto fail;
+            }
+        }
         return res;
     }
     else if (whence != 0) {