Make test_marshal pass. Not my best work. :-(

diff --git a/Lib/test/test_marshal.py b/Lib/test/test_marshal.py
index 7f7c5e6cdd688d065696ed4cb1a61696a5222635..2af6ae7c3a8d0e66546acac70cdb89e4233b0555 100644 (file)
--- a/Lib/test/test_marshal.py
+++ b/Lib/test/test_marshal.py
@@ -27,18 +27,18 @@ class IntTestCase(unittest.TestCase):
         # we're running the test on a 32-bit box, of course.
 
         def to_little_endian_string(value, nbytes):
-            bytes = []
+            b = bytes()
             for i in range(nbytes):
-                bytes.append(chr(value & 0xff))
+                b.append(value & 0xff)
                 value >>= 8
-            return ''.join(bytes)
+            return b
 
         maxint64 = (1 << 63) - 1
         minint64 = -maxint64-1
 
         for base in maxint64, minint64, -maxint64, -(minint64 >> 1):
             while base:
-                s = 'I' + to_little_endian_string(base, 8)
+                s = b'I' + to_little_endian_string(base, 8)
                 got = marshal.loads(s)
                 self.assertEqual(base, got)
                 if base == -1:  # a fixed-point for shifting right 1
@@ -128,7 +128,7 @@ class StringTestCase(unittest.TestCase):
         os.unlink(test_support.TESTFN)
 
     def test_buffer(self):
-        for s in ["", "Andrè Previn", "abc", " "*10000]:
+        for s in [b"", b"Andr\xe8 Previn", b"abc", b" "*10000]:
             b = buffer(s)
             new = marshal.loads(marshal.dumps(b))
             self.assertEqual(s, new)

diff --git a/Python/marshal.c b/Python/marshal.c
index 9243798e4bbeaf76efe42d0e68b717fd91d48a8b..72253eaccdeaeb38b217d546f9a48a20021a6a88 100644 (file)
--- a/Python/marshal.c
+++ b/Python/marshal.c
@@ -1014,6 +1014,8 @@ PyObject *
 PyMarshal_WriteObjectToString(PyObject *x, int version)
 {
        WFILE wf;
+       PyObject *res = NULL;
+
        wf.fp = NULL;
        wf.str = PyString_FromStringAndSize((char *)NULL, 50);
        if (wf.str == NULL)
@@ -1034,7 +1036,8 @@ PyMarshal_WriteObjectToString(PyObject *x, int version)
                                        "too much marshal data for a string");
                        return NULL;
                }
-               _PyString_Resize(&wf.str, (Py_ssize_t)(wf.ptr - base));
+               if (_PyString_Resize(&wf.str, (Py_ssize_t)(wf.ptr - base)) < 0)
+                       return NULL;
        }
        if (wf.error) {
                Py_XDECREF(wf.str);
@@ -1043,7 +1046,12 @@ PyMarshal_WriteObjectToString(PyObject *x, int version)
                                :"object too deeply nested to marshal");
                return NULL;
        }
-       return wf.str;
+       if (wf.str != NULL) {
+               /* XXX Quick hack -- need to do this differently */
+               res = PyBytes_FromObject(wf.str);
+               Py_DECREF(wf.str);
+       }
+       return res;
 }
 
 /* And an interface for Python programs... */
@@ -1092,9 +1100,34 @@ marshal_load(PyObject *self, PyObject *f)
        RFILE rf;
        PyObject *result;
        if (!PyFile_Check(f)) {
-               PyErr_SetString(PyExc_TypeError,
-                               "marshal.load() arg must be file");
-               return NULL;
+               /* XXX Quick hack -- need to do this differently */
+               PyObject *data, *result;
+               RFILE rf;
+               data = PyObject_CallMethod(f, "read", "");
+               if (data == NULL)
+                       return NULL;
+               rf.fp = NULL;
+               if (PyString_Check(data)) {
+                       rf.ptr = PyString_AS_STRING(data);
+                       rf.end = rf.ptr + PyString_GET_SIZE(data);
+               }
+               else if (PyBytes_Check(data)) {
+                       rf.ptr = PyBytes_AS_STRING(data);
+                       rf.end = rf.ptr + PyBytes_GET_SIZE(data);
+               }
+               else {
+                       PyErr_Format(PyExc_TypeError,
+                                    "f.read() returned neither string "
+                                    "nor bytes but %.100s",
+                                    data->ob_type->tp_name);
+                       Py_DECREF(data);
+                       return NULL;
+               }
+               rf.strings = PyList_New(0);
+               result = read_object(&rf);
+               Py_DECREF(rf.strings);
+               Py_DECREF(data);
+               return result;
        }
        rf.fp = PyFile_AsFile(f);
        rf.strings = PyList_New(0);