diff --git a/Lib/test/test_multibytecodec.py b/Lib/test/test_multibytecodec.py --- a/Lib/test/test_multibytecodec.py +++ b/Lib/test/test_multibytecodec.py @@ -6,7 +6,7 @@ from test import support from test.support import TESTFN -import unittest, io, codecs, sys, os +import unittest, io, codecs, sys, os, tempfile import _multibytecodec ALL_CJKENCODINGS = [ @@ -256,6 +256,52 @@ class Test_ISO2022(unittest.TestCase): # Any ISO 2022 codec will cause the segfault myunichr(x).encode('iso_2022_jp', 'ignore') +class TestStateful(unittest.TestCase): + text = '\u4E16\u4E16' + encoding = 'iso-2022-jp' + expected = b'\x1b$B@$@$' + expected_reset = b'\x1b$B@$@$\x1b(B' + + def test_encode(self): + self.assertEqual(self.text.encode(self.encoding), self.expected_reset) + + def test_incrementalencoder(self): + encoder = codecs.getincrementalencoder(self.encoding)() + output = b''.join( + encoder.encode(char) + for char in self.text) + self.assertEqual(output, self.expected) + + def test_incrementalencoder_final(self): + encoder = codecs.getincrementalencoder(self.encoding)() + last_index = len(self.text) - 1 + output = b''.join( + encoder.encode(char, index == last_index) + for index, char in enumerate(self.text)) + self.assertEqual(output, self.expected_reset) + + def test_file_write(self): + with tempfile.NamedTemporaryFile() as temp: + with open(temp.name, "w", encoding=self.encoding) as f: + f.write(self.text) + with open(temp.name, "rb") as f: + output = f.read() + self.assertEqual(output, self.expected) + + def test_file_writelines(self): + with tempfile.NamedTemporaryFile() as temp: + with open(temp.name, "w", encoding=self.encoding) as f: + f.writelines(self.text) + with open(temp.name, "rb") as f: + output = f.read() + self.assertEqual(output, self.expected) + +class TestHZStateful(TestStateful): + text = '\u804a\u804a' + encoding = 'hz' + expected = b'~{ADAD' + expected_reset = b'~{ADAD~}' + def test_main(): support.run_unittest(__name__) diff --git a/Modules/cjkcodecs/multibytecodec.c b/Modules/cjkcodecs/multibytecodec.c --- a/Modules/cjkcodecs/multibytecodec.c +++ b/Modules/cjkcodecs/multibytecodec.c @@ -479,7 +479,7 @@ multibytecodec_encode(MultibyteCodec *co MultibyteEncodeBuffer buf; Py_ssize_t finalsize, r = 0; - if (datalen == 0) + if (datalen == 0 && !(flags & MBENC_RESET)) return PyBytes_FromStringAndSize(NULL, 0); buf.excobj = NULL; @@ -515,7 +515,7 @@ multibytecodec_encode(MultibyteCodec *co break; } - if (codec->encreset != NULL) + if (codec->encreset != NULL && (flags & MBENC_RESET)) for (;;) { Py_ssize_t outleft;