Restore Python 2.1 StringIO.py behaviour: support concatenating

malemburg · malemburg · commit f853be980ee5 · 2002-01-06T17:15:05.000Z
Unicode string snippets to larger Unicode strings.

This fix should also go into Python 2.2.1.
diff --git a/Lib/StringIO.py b/Lib/StringIO.py
@@ -28,7 +28,7 @@
   bytes that occupy space in the buffer.
 - There's a simple test set (see end of this file).
 """
-
+import types
 try:
     from errno import EINVAL
 except ImportError:
@@ -38,8 +38,10 @@
 
 class StringIO:
     def __init__(self, buf = ''):
-        # Force self.buf to be a string
-        self.buf = str(buf)
+        # Force self.buf to be a string or unicode
+        if type(buf) is not types.UnicodeType:
+            buf = str(buf)
+        self.buf = buf
         self.len = len(buf)
         self.buflist = []
         self.pos = 0
@@ -135,8 +137,9 @@ def write(self, s):
         if self.closed:
             raise ValueError, "I/O operation on closed file"
         if not s: return
-        # Force s to be a string
-        s = str(s)
+        # Force s to be a string or unicode
+        if type(s) is not types.UnicodeType:
+            s = str(s)
         if self.pos > self.len:
             self.buflist.append('\0'*(self.pos - self.len))
             self.len = self.pos
diff --git a/Lib/test/test_StringIO.py b/Lib/test/test_StringIO.py
@@ -71,6 +71,21 @@ def test_iterator(self):
 class TestStringIO(TestGenericStringIO):
     MODULE = StringIO
 
+    def test_unicode(self):
+
+        # The StringIO module also supports concatenating Unicode
+        # snippets to larger Unicode strings. This is tested by this
+        # method. Note that cStringIO does not support this extension.
+        
+        f = self.MODULE.StringIO()
+        f.write(self._line[:6])
+        f.seek(3)
+        f.write(unicode(self._line[20:26]))
+        f.write(unicode(self._line[52]))
+        s = f.getvalue()
+        self.assertEqual(s, unicode('abcuvwxyz!'))
+        self.assertEqual(type(s), types.UnicodeType)
+
 class TestcStringIO(TestGenericStringIO):
     MODULE = cStringIO
 
