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

gvanrossum · gvanrossum · commit e6d3904c224b · 2007-05-09T00:01:30.000Z
diff --git a/Lib/test/test_marshal.py b/Lib/test/test_marshal.py
@@ -27,18 +27,18 @@ def test_int64(self):
         # 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 @@ def test_string(self):
         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
@@ -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);
