Skip to content

Commit d86d637

Browse files
committed
gh-160185: Fix CJK StreamWriter.reset() not returning to ASCII
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 ab5187c commit d86d637

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
@@ -438,6 +438,42 @@ def test_incrementalencoder_final(self):
438438
self.assertEqual(output, self.expected_reset)
439439
self.assertEqual(encoder.encode('', final=True), b'')
440440

441+
def test_streamwriter_reset(self):
442+
stream = io.BytesIO()
443+
writer = codecs.getwriter(self.encoding)(stream)
444+
writer.reset()
445+
self.assertEqual(stream.getvalue(), b'')
446+
writer.write(self.text)
447+
self.assertEqual(stream.getvalue(), self.expected)
448+
writer.reset()
449+
self.assertEqual(stream.getvalue(), self.expected_reset)
450+
writer.reset()
451+
self.assertEqual(stream.getvalue(), self.expected_reset)
452+
writer.write(self.text)
453+
writer.reset()
454+
self.assertEqual(stream.getvalue(), self.expected_reset * 2)
455+
456+
def test_streamwriter_seek0(self):
457+
# The reset sequence is written before the stream is rewound.
458+
info = codecs.lookup(self.encoding)
459+
for wrap in (
460+
info.streamwriter,
461+
lambda stream: codecs.StreamReaderWriter(
462+
stream, info.streamreader, info.streamwriter),
463+
lambda stream: codecs.EncodedFile(stream, 'utf-8', self.encoding),
464+
):
465+
stream = io.BytesIO()
466+
writer = wrap(stream)
467+
text = self.text
468+
if isinstance(writer, codecs.StreamRecoder):
469+
text = text.encode()
470+
writer.write(text)
471+
writer.seek(0)
472+
self.assertEqual(stream.tell(), 0)
473+
self.assertEqual(stream.getvalue(), self.expected_reset)
474+
writer.write(text)
475+
self.assertEqual(stream.getvalue(), self.expected_reset)
476+
441477
class TestHZStateful(TestStateful):
442478
text = '\u804a\u804a'
443479
encoding = 'hz'
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
@@ -1857,13 +1857,17 @@ _multibytecodec_MultibyteStreamWriter_reset_impl(MultibyteStreamWriterObject *se
18571857
PyTypeObject *cls)
18581858
/*[clinic end generated code: output=32ef224c2a38aa3d input=28af6a9cd38d1979]*/
18591859
{
1860-
PyObject *pwrt;
1860+
PyObject *pwrt, *pending;
18611861

1862-
if (!self->pending)
1863-
Py_RETURN_NONE;
1862+
/* even without pending characters, a stateful codec may have to
1863+
* write a sequence that returns the output to its initial state */
1864+
pending = self->pending;
1865+
if (pending == NULL) {
1866+
pending = Py_GetConstantBorrowed(Py_CONSTANT_EMPTY_STR);
1867+
}
18641868

18651869
pwrt = multibytecodec_encode(self->codec, &self->state,
1866-
self->pending, NULL, self->errors,
1870+
pending, NULL, self->errors,
18671871
MBENC_FLUSH | MBENC_RESET);
18681872
/* some pending buffer can be truncated when UnicodeEncodeError is
18691873
* raised on 'strict' mode. but, 'reset' method is designed to
@@ -1886,6 +1890,7 @@ _multibytecodec_MultibyteStreamWriter_reset_impl(MultibyteStreamWriterObject *se
18861890
Py_DECREF(pwrt);
18871891
return NULL;
18881892
}
1893+
Py_DECREF(wr);
18891894
}
18901895
Py_DECREF(pwrt);
18911896

0 commit comments

Comments
 (0)