StringIO patch #462596: let's [c]StringIO accept read buffers on

input to .write() too.

3 files changed, 32 insertions, 18 deletions

diff --git a/Lib/StringIO.py b/Lib/StringIO.py
index cb7313d..e7c8e04 100644
--- a/Lib/StringIO.py
+++ b/Lib/StringIO.py
@@ -38,7 +38,8 @@ __all__ = ["StringIO"]
 
 class StringIO:
     def __init__(self, buf = ''):
-        self.buf = buf
+        # Force self.buf to be a string
+        self.buf = str(buf)
         self.len = len(buf)
         self.buflist = []
         self.pos = 0
@@ -134,6 +135,8 @@ class StringIO:
         if self.closed:
             raise ValueError, "I/O operation on closed file"
         if not s: return
+        # Force s to be a string
+        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
index eddeb10..4a0a814 100644
--- a/Lib/test/test_StringIO.py
+++ b/Lib/test/test_StringIO.py
@@ -10,28 +10,35 @@ import test_support
 class TestGenericStringIO(unittest.TestCase):
     # use a class variable MODULE to define which module is being tested
 
+    # Line of data to test as string
+    _line = 'abcdefghijklmnopqrstuvwxyzABCDEFGHIJKLMNOPQRSTUVWXYZ!'
+
+    # Constructor to use for the test data (._line is passed to this
+    # constructor)
+    constructor = str
+
     def setUp(self):
-        self._line = 'abcdefghijklmnopqrstuvwxyzABCDEFGHIJKLMNOPQRSTUVWXYZ'
-        self._lines = (self._line + '\n') * 5
+        self._line = self.constructor(self._line)
+        self._lines = self.constructor((self._line + '\n') * 5)
         self._fp = self.MODULE.StringIO(self._lines)
 
     def test_reads(self):
         eq = self.assertEqual
-        eq(self._fp.read(10), 'abcdefghij')
-        eq(self._fp.readline(), 'klmnopqrstuvwxyzABCDEFGHIJKLMNOPQRSTUVWXYZ\n')
+        eq(self._fp.read(10), self._line[:10])
+        eq(self._fp.readline(), self._line[10:] + '\n')
         eq(len(self._fp.readlines(60)), 2)
 
     def test_writes(self):
         f = self.MODULE.StringIO()
-        f.write('abcdef')
+        f.write(self._line[:6])
         f.seek(3)
-        f.write('uvwxyz')
-        f.write('!')
+        f.write(self._line[20:26])
+        f.write(self._line[52])
         self.assertEqual(f.getvalue(), 'abcuvwxyz!')
 
     def test_writelines(self):
         f = self.MODULE.StringIO()
-        f.writelines(['a', 'b', 'c'])
+        f.writelines([self._line[0], self._line[1], self._line[2]])
         f.seek(0)
         self.assertEqual(f.getvalue(), 'abc')
 
@@ -64,10 +71,18 @@ class TestStringIO(TestGenericStringIO):
 class TestcStringIO(TestGenericStringIO):
     MODULE = cStringIO
 
+class TestBufferStringIO(TestStringIO):
+    constructor = buffer
+
+class TestBuffercStringIO(TestcStringIO):
+    constructor = buffer
+
 
 def test_main():
     test_support.run_unittest(TestStringIO)
     test_support.run_unittest(TestcStringIO)
+    test_support.run_unittest(TestBufferStringIO)
+    test_support.run_unittest(TestBuffercStringIO)
 
 if __name__ == '__main__':
     test_main()
diff --git a/Modules/cStringIO.c b/Modules/cStringIO.c
index 36ed354..49007e2 100644
--- a/Modules/cStringIO.c
+++ b/Modules/cStringIO.c
@@ -120,7 +120,8 @@ typedef struct { /* Subtype of IOobject */
   PyObject_HEAD
   char *buf;
   int pos, string_size;
-
+  /* We store a reference to the object here in order to keep
+     the buffer alive during the lifetime of the Iobject. */
   PyObject *pbuf;
 } Iobject;
 
@@ -424,14 +425,11 @@ O_cwrite(PyObject *self, char *c, int  l) {
 
 static PyObject *
 O_write(Oobject *self, PyObject *args) {
-        PyObject *s;
         char *c;
         int l;
 
-        UNLESS (PyArg_ParseTuple(args, "O:write", &s)) return NULL;
+        UNLESS (PyArg_ParseTuple(args, "s#:write", &c, &l)) return NULL;
 
-        UNLESS (-1 != (l=PyString_Size(s))) return NULL;
-        UNLESS (c=PyString_AsString(s)) return NULL;
         if (O_cwrite((PyObject*)self,c,l) < 0) return NULL;
 
         Py_INCREF(Py_None);
@@ -713,13 +711,11 @@ newIobject(PyObject *s) {
   char *buf;
   int size;
 
-  if (!PyString_Check(s)) {
-      PyErr_Format(PyExc_TypeError, "expected string, %.200s found",
+  if (PyObject_AsReadBuffer(s, (const void **)&buf, &size)) {
+      PyErr_Format(PyExc_TypeError, "expected read buffer, %.200s found",
 		   s->ob_type->tp_name);
       return NULL;
   }
-  buf = PyString_AS_STRING(s);
-  size = PyString_GET_SIZE(s);
   UNLESS (self = PyObject_New(Iobject, &Itype)) return NULL;
   Py_INCREF(s);
   self->buf=buf;