Skip to content

Commit aa4f541

Browse files
authored
gh-160185: Fix CJK StreamWriter.reset() not returning to ASCII (GH-160186)
reset() returned early when nothing was pending, so the shift-back sequence of hz and iso2022_* was never written (regression from gh-67436). Also fix a reference leak of the stream.write() result.
1 parent 5f28e31 commit aa4f541

4 files changed

Lines changed: 56 additions & 8 deletions

File tree

‎Lib/codecs.py‎

Lines changed: 6 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -402,9 +402,10 @@ def reset(self):
402402
pass
403403

404404
def seek(self, offset, whence=0):
405-
self.stream.seek(offset, whence)
405+
# reset() may write to the stream: call it at the old position.
406406
if whence == 0 and offset == 0:
407407
self.reset()
408+
self.stream.seek(offset, whence)
408409

409410
def __getattr__(self, name,
410411
getattr=getattr):
@@ -741,10 +742,11 @@ def reset(self):
741742
self.writer.reset()
742743

743744
def seek(self, offset, whence=0):
744-
self.stream.seek(offset, whence)
745-
self.reader.reset()
745+
# writer.reset() may write to the stream: call it at the old position.
746746
if whence == 0 and offset == 0:
747747
self.writer.reset()
748+
self.stream.seek(offset, whence)
749+
self.reader.reset()
748750

749751
def __getattr__(self, name,
750752
getattr=getattr):
@@ -864,8 +866,8 @@ def reset(self):
864866
def seek(self, offset, whence=0):
865867
# Seeks must be propagated to both the readers and writers
866868
# as they might need to reset their internal buffers.
867-
self.reader.seek(offset, whence)
868869
self.writer.seek(offset, whence)
870+
self.reader.seek(offset, whence)
869871

870872
def __getattr__(self, name,
871873
getattr=getattr):

‎Lib/test/test_multibytecodec.py‎

Lines changed: 36 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -461,6 +461,42 @@ def test_incrementalencoder_final(self):
461461
self.assertEqual(output, self.expected_reset)
462462
self.assertEqual(encoder.encode('', final=True), b'')
463463

464+
def test_streamwriter_reset(self):
465+
stream = io.BytesIO()
466+
writer = codecs.getwriter(self.encoding)(stream)
467+
writer.reset()
468+
self.assertEqual(stream.getvalue(), b'')
469+
writer.write(self.text)
470+
self.assertEqual(stream.getvalue(), self.expected)
471+
writer.reset()
472+
self.assertEqual(stream.getvalue(), self.expected_reset)
473+
writer.reset()
474+
self.assertEqual(stream.getvalue(), self.expected_reset)
475+
writer.write(self.text)
476+
writer.reset()
477+
self.assertEqual(stream.getvalue(), self.expected_reset * 2)
478+
479+
def test_streamwriter_seek0(self):
480+
# The reset sequence is written before the stream is rewound.
481+
info = codecs.lookup(self.encoding)
482+
for wrap in (
483+
info.streamwriter,
484+
lambda stream: codecs.StreamReaderWriter(
485+
stream, info.streamreader, info.streamwriter),
486+
lambda stream: codecs.EncodedFile(stream, 'utf-8', self.encoding),
487+
):
488+
stream = io.BytesIO()
489+
writer = wrap(stream)
490+
text = self.text
491+
if isinstance(writer, codecs.StreamRecoder):
492+
text = text.encode()
493+
writer.write(text)
494+
writer.seek(0)
495+
self.assertEqual(stream.tell(), 0)
496+
self.assertEqual(stream.getvalue(), self.expected_reset)
497+
writer.write(text)
498+
self.assertEqual(stream.getvalue(), self.expected_reset)
499+
464500
def test_incrementalencoder_error(self):
465501
# The output of a failed encode() is lost, so the shift state
466502
# must not be changed.
Lines changed: 5 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,5 @@
1+
Fix :meth:`codecs.StreamWriter.reset` of the stateful CJK codecs (``hz`` and
2+
the ``iso2022_*`` family) not writing the sequence that returns the output to
3+
ASCII. ``seek(0)`` of :class:`codecs.StreamWriter`,
4+
:class:`codecs.StreamReaderWriter` and :class:`codecs.StreamRecoder` now
5+
resets the writer before moving the stream position.

‎Modules/cjkcodecs/multibytecodec.c‎

Lines changed: 9 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -1875,13 +1875,17 @@ _multibytecodec_MultibyteStreamWriter_reset_impl(MultibyteStreamWriterObject *se
18751875
PyTypeObject *cls)
18761876
/*[clinic end generated code: output=32ef224c2a38aa3d input=28af6a9cd38d1979]*/
18771877
{
1878-
PyObject *pwrt;
1878+
PyObject *pwrt, *pending;
18791879

1880-
if (!self->pending)
1881-
Py_RETURN_NONE;
1880+
/* even without pending characters, a stateful codec may have to
1881+
* write a sequence that returns the output to its initial state */
1882+
pending = self->pending;
1883+
if (pending == NULL) {
1884+
pending = Py_GetConstantBorrowed(Py_CONSTANT_EMPTY_STR);
1885+
}
18821886

18831887
pwrt = multibytecodec_encode(self->codec, &self->state,
1884-
self->pending, NULL, self->errors,
1888+
pending, NULL, self->errors,
18851889
MBENC_FLUSH | MBENC_RESET);
18861890
/* some pending buffer can be truncated when UnicodeEncodeError is
18871891
* raised on 'strict' mode. but, 'reset' method is designed to
@@ -1904,6 +1908,7 @@ _multibytecodec_MultibyteStreamWriter_reset_impl(MultibyteStreamWriterObject *se
19041908
Py_DECREF(pwrt);
19051909
return NULL;
19061910
}
1911+
Py_DECREF(wr);
19071912
}
19081913
Py_DECREF(pwrt);
19091914

0 commit comments

Comments
 (0)