diff options
Diffstat (limited to 'Lib')
| -rw-r--r-- | Lib/codecs.py | 33 | ||||
| -rw-r--r-- | Lib/test/test_codecs.py | 29 |
2 files changed, 58 insertions, 4 deletions
diff --git a/Lib/codecs.py b/Lib/codecs.py index 1518d75..f834b8d 100644 --- a/Lib/codecs.py +++ b/Lib/codecs.py @@ -329,6 +329,12 @@ class StreamWriter(Codec): """ return getattr(self.stream, name) + def __enter__(self): + return self + + def __exit__(self, type, value, tb): + self.stream.close() + ### class StreamReader(Codec): @@ -568,6 +574,12 @@ class StreamReader(Codec): """ return getattr(self.stream, name) + def __enter__(self): + return self + + def __exit__(self, type, value, tb): + self.stream.close() + ### class StreamReaderWriter: @@ -641,6 +653,14 @@ class StreamReaderWriter: """ return getattr(self.stream, name) + # these are needed to make "with codecs.open(...)" work properly + + def __enter__(self): + return self + + def __exit__(self, type, value, tb): + self.stream.close() + ### class StreamRecoder: @@ -751,6 +771,12 @@ class StreamRecoder: """ return getattr(self.stream, name) + def __enter__(self): + return self + + def __exit__(self, type, value, tb): + self.stream.close() + ### Shortcuts def open(filename, mode='rb', encoding=None, errors='strict', buffering=1): @@ -824,9 +850,10 @@ def EncodedFile(file, data_encoding, file_encoding=None, errors='strict'): """ if file_encoding is None: file_encoding = data_encoding - info = lookup(data_encoding) - sr = StreamRecoder(file, info.encode, info.decode, - info.streamreader, info.streamwriter, errors) + data_info = lookup(data_encoding) + file_info = lookup(file_encoding) + sr = StreamRecoder(file, data_info.encode, data_info.decode, + file_info.streamreader, file_info.streamwriter, errors) # Add attributes to simplify introspection sr.data_encoding = data_encoding sr.file_encoding = file_encoding diff --git a/Lib/test/test_codecs.py b/Lib/test/test_codecs.py index 8153979..39d4206 100644 --- a/Lib/test/test_codecs.py +++ b/Lib/test/test_codecs.py @@ -910,6 +910,18 @@ class StreamReaderTest(unittest.TestCase): f = self.reader(self.stream) self.assertEquals(f.readlines(), [u'\ud55c\n', u'\uae00']) +class EncodedFileTest(unittest.TestCase): + + def test_basic(self): + f = StringIO.StringIO('\xed\x95\x9c\n\xea\xb8\x80') + ef = codecs.EncodedFile(f, 'utf-16', 'utf-8') + self.assertEquals(ef.read(), '\xff\xfe\\\xd5\n\x00\x00\xae') + + f = StringIO.StringIO() + ef = codecs.EncodedFile(f, 'utf-8', 'latin1') + ef.write('\xc3\xbc') + self.assertEquals(f.getvalue(), '\xfc') + class Str2StrTest(unittest.TestCase): def test_read(self): @@ -1214,6 +1226,19 @@ class CharmapTest(unittest.TestCase): (u"", len(allbytes)) ) +class WithStmtTest(unittest.TestCase): + def test_encodedfile(self): + f = StringIO.StringIO("\xc3\xbc") + with codecs.EncodedFile(f, "latin-1", "utf-8") as ef: + self.assertEquals(ef.read(), "\xfc") + + def test_streamreaderwriter(self): + f = StringIO.StringIO("\xc3\xbc") + info = codecs.lookup("utf-8") + with codecs.StreamReaderWriter(f, info.streamreader, + info.streamwriter, 'strict') as srw: + self.assertEquals(srw.read(), u"\xfc") + def test_main(): test_support.run_unittest( @@ -1234,10 +1259,12 @@ def test_main(): IDNACodecTest, CodecsModuleTest, StreamReaderTest, + EncodedFileTest, Str2StrTest, BasicUnicodeTest, BasicStrTest, - CharmapTest + CharmapTest, + WithStmtTest, ) |