[Python-checkins] r57040 - in python/branches/alex-py3k: Lib/test/test_memoryio.py Modules/_bytesiomodule.c

[alexandre.vassalotti]

Author: alexandre.vassalotti
Date: Tue Aug 14 23:43:19 2007
New Revision: 57040

Modified:
   python/branches/alex-py3k/Lib/test/test_memoryio.py
   python/branches/alex-py3k/Modules/_bytesiomodule.c
Log:
Encode Unicode strings given to BytesIO objects with UTF-8 before
storing them as bytes.


Modified: python/branches/alex-py3k/Lib/test/test_memoryio.py
==============================================================================
--- python/branches/alex-py3k/Lib/test/test_memoryio.py	(original)
+++ python/branches/alex-py3k/Lib/test/test_memoryio.py	Tue Aug 14 23:43:19 2007
@@ -196,7 +196,7 @@
         self.assertEqual(memio.isatty(), False)
         self.assertEqual(memio.closed, False)
         memio.close()
-        self.assertEqual(memio.writable(), True)
+        self.assertEqual(memio.writable(), True) # closed
         self.assertEqual(memio.readable(), True)
         self.assertEqual(memio.seekable(), True)
         self.assertRaises(ValueError, memio.isatty)
@@ -211,6 +211,16 @@
             return m.getvalue()
         self.assertEqual(test(), buf)
 
+    def test_widechar(self):
+        buf = self.buftype("\U0002030a\U00020347")
+        memio = self.ioclass(buf)
+
+        self.assertEqual(memio.getvalue(), buf)
+        self.assertEqual(memio.write(buf), len(buf))
+        self.assertEqual(memio.getvalue(), buf)
+        self.assertEqual(memio.write(buf), len(buf))
+        self.assertEqual(memio.getvalue(), buf + buf)
+
 
 class PyBytesIOTest(MemoryTestMixin, unittest.TestCase):
     buftype = bytes
@@ -247,6 +257,16 @@
         memio.write(buf)
         self.assertEqual(memio.getvalue(), self.buftype('\0\0') + buf)
 
+    def test_unicode(self):
+        buf = "1234567890"
+        memio = self.ioclass(buf)
+
+        self.assertEqual(memio.getvalue(), self.buftype(buf))
+        self.assertEqual(memio.write(buf), len(buf))
+        self.assertEqual(memio.getvalue(), self.buftype(buf))
+        self.assertEqual(memio.write(buf), len(buf))
+        self.assertEqual(memio.getvalue(), self.buftype(buf + buf))
+
 
 class PyStringIOTest(MemoryTestMixin, unittest.TestCase):
     buftype = str

Modified: python/branches/alex-py3k/Modules/_bytesiomodule.c
==============================================================================
--- python/branches/alex-py3k/Modules/_bytesiomodule.c	(original)
+++ python/branches/alex-py3k/Modules/_bytesiomodule.c	Tue Aug 14 23:43:19 2007
@@ -395,16 +395,22 @@
 static PyObject *
 bytesio_write(BytesIOObject *self, PyObject *obj)
 {
-    const void *bytes;
+    const char *bytes;
     Py_ssize_t size, len;
 
     if (self->buf == NULL)
         return err_closed();
 
-    if (PyObject_AsReadBuffer(obj, &bytes, &size) < 0)
-        return NULL;
+    if (PyUnicode_Check(obj)) {
+        bytes = PyUnicode_AsString(obj);
+        size = strlen(bytes);
+    }
+    else {
+        if (PyObject_AsReadBuffer(obj, (void *)&bytes, &size) < 0)
+            return NULL;
+    }
 
-    len = write_bytes(self, (const char *)bytes, size);
+    len = write_bytes(self, bytes, size);
     if (len == -1)
         return NULL;
 
@@ -466,12 +472,12 @@
 BytesIO_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
 {
     BytesIOObject *self;
-    const char *initvalue = NULL;
-    Py_ssize_t size = 0;
+    PyObject *initvalue = NULL, *ret;
+    enum { INIT_BUFSIZE = 1 };
 
     assert(type != NULL && type->tp_alloc != NULL);
 
-    if (!PyArg_ParseTuple(args, "|s#ss:BytesIO", &initvalue, &size))
+    if (!PyArg_ParseTuple(args, "|O:BytesIO", &initvalue))
         return NULL;
 
     self = (BytesIOObject *)type->tp_alloc(type, 0);
@@ -479,25 +485,21 @@
     if (self == NULL)
         return NULL;
 
-    self->buf = PyMem_New(char, size + 1);
+    self->buf = PyMem_New(char, INIT_BUFSIZE);
 
     /* These variables need to be initialized before attempting to write
        anything to the object. */
     self->pos = 0;
     self->string_size = 0;
-    self->buf_size = size + 1;
+    self->buf_size = INIT_BUFSIZE;
 
-    if (size > 0) {
-        if (write_bytes(self, initvalue, size) == -1) {
-            Py_DECREF(self);            
+    if (initvalue) {
+        ret = bytesio_write(self, initvalue);
+        if (ret == NULL)
             return NULL;
-        }
+        Py_DECREF(ret);
         self->pos = 0;
     }
-    if (self->buf == NULL) {
-        Py_DECREF(self);
-        return PyErr_NoMemory();
-    }
 
     return (PyObject *)self;
 }