From 2a5a3027f23a2c03e20e1c382be08ebb390f3ea7 Mon Sep 17 00:00:00 2001
From: Georg Brandl <georg@python.org>
Date: Sun, 29 Oct 2006 08:39:27 +0000
Subject: [PATCH] Fix codecs.EncodedFile which did not use file_encoding in
 2.5.0, and fix all codecs file wrappers to work correctly with the "with"
 statement (bug #1586513).  (backport from rev. 52517)

---
 Lib/codecs.py           | 33 ++++++++++++++++++++++++++++++---
 Lib/test/test_codecs.py | 30 +++++++++++++++++++++++++++++-
 Misc/NEWS               |  4 ++++
 3 files changed, 63 insertions(+), 4 deletions(-)

diff --git a/Lib/codecs.py b/Lib/codecs.py
index 1518d75f9d..f834b8dd17 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 8153979a85..90340bbcb0 100644
--- a/Lib/test/test_codecs.py
+++ b/Lib/test/test_codecs.py
@@ -1,3 +1,4 @@
+from __future__ import with_statement
 from test import test_support
 import unittest
 import codecs
@@ -910,6 +911,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 +1227,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 +1260,12 @@ def test_main():
         IDNACodecTest,
         CodecsModuleTest,
         StreamReaderTest,
+        EncodedFileTest,
         Str2StrTest,
         BasicUnicodeTest,
         BasicStrTest,
-        CharmapTest
+        CharmapTest,
+        WithStmtTest,
     )
 
 
diff --git a/Misc/NEWS b/Misc/NEWS
index f47c4e920e..b8a6f7e392 100644
--- a/Misc/NEWS
+++ b/Misc/NEWS
@@ -90,6 +90,10 @@ Extension Modules
 Library
 -------
 
+- Fix codecs.EncodedFile which did not use file_encoding in 2.5.0, and
+  fix all codecs file wrappers to work correctly with the "with"
+  statement (bug #1586513).
+
 - ctypes callback functions only support 'fundamental' data types as
   result type.  Raise an error when something else is used.  This is a
   partial fix for Bug #1574584.
-- 
2.40.0