Skip to content
Open
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
42 changes: 42 additions & 0 deletions Lib/test/test_multibytecodec.py
Original file line number Diff line number Diff line change
Expand Up @@ -438,6 +438,48 @@ def test_incrementalencoder_final(self):
self.assertEqual(output, self.expected_reset)
self.assertEqual(encoder.encode('', final=True), b'')

def test_incrementalencoder_error(self):
# The output of a failed encode() is lost, so the shift state
# must not be changed.
encoder = codecs.getincrementalencoder(self.encoding)()
state = encoder.getstate()
for final in False, True:
self.assertRaises(UnicodeEncodeError, encoder.encode,
self.text + '\U0001f600', final)
self.assertEqual(encoder.getstate(), state)
self.assertEqual(encoder.encode(self.text), self.expected)
state = encoder.getstate()
for final in False, True:
self.assertRaises(UnicodeEncodeError, encoder.encode,
'a\U0001f600', final)
self.assertEqual(encoder.getstate(), state)
self.assertEqual(encoder.encode('a'), self.reset + b'a')

def test_streamwriter_error(self):
stream = io.BytesIO()
writer = codecs.getwriter(self.encoding)(stream)
self.assertRaises(UnicodeEncodeError, writer.write,
self.text + '\U0001f600')
self.assertEqual(stream.getvalue(), b'')
writer.write(self.text)
self.assertEqual(stream.getvalue(), self.expected)

def test_textiowrapper_error(self):
raw = io.BytesIO()
with io.TextIOWrapper(raw, encoding=self.encoding) as f:
self.assertRaises(UnicodeEncodeError, f.write,
self.text + '\U0001f600')
f.write(self.text)
f.flush()
self.assertEqual(raw.getvalue(), self.expected)

class TestISO2022KRStateful(TestStateful):
text = '\uac00\uac00'
encoding = 'iso-2022-kr'
expected = b'\x1b$)C\x0e0!0!'
reset = b'\x0f'
expected_reset = expected + reset

class TestHZStateful(TestStateful):
text = '\u804a\u804a'
encoding = 'hz'
Expand Down
Original file line number Diff line number Diff line change
@@ -0,0 +1,4 @@
Fix the incremental encoders and stream writers of the stateful CJK codecs
(``iso2022_*`` and ``hz``) keeping a changed shift state after an
:meth:`!encode` or :meth:`!write` call failed. The next call could omit the
escape sequence, writing multibyte characters in ASCII mode.
5 changes: 5 additions & 0 deletions Modules/cjkcodecs/multibytecodec.c
Original file line number Diff line number Diff line change
Expand Up @@ -787,6 +787,7 @@ encoder_encode_stateful(MultibyteStatefulEncoderContext *ctx,
PyObject *inbuf = NULL;
Py_ssize_t inpos, datalen;
PyObject *origpending = NULL;
MultibyteCodec_State origstate;

if (PyUnicode_Check(unistr))
ucvt = NULL;
Expand All @@ -802,6 +803,7 @@ encoder_encode_stateful(MultibyteStatefulEncoderContext *ctx,
}
}

origstate = ctx->state;
if (ctx->pending) {
PyObject *inbuf_tmp;

Expand Down Expand Up @@ -859,6 +861,9 @@ encoder_encode_stateful(MultibyteStatefulEncoderContext *ctx,
return r;

errorexit:
/* recover the original codec state: the discarded output may have
* included an escape sequence for a state change */
ctx->state = origstate;
Py_XDECREF(r);
Py_XDECREF(ucvt);
Py_XDECREF(origpending);
Expand Down
Loading