SF 1098985: set objects cannot be marshalled

3 files changed, 80 insertions, 1 deletions

diff --git a/Lib/test/test_marshal.py b/Lib/test/test_marshal.py
index 5c1a3f3..9901a3c 100644
--- a/Lib/test/test_marshal.py
+++ b/Lib/test/test_marshal.py
@@ -170,6 +170,18 @@ class ContainerTestCase(unittest.TestCase):
         self.assertEqual(t, new)
         os.unlink(test_support.TESTFN)
 
+    def test_sets(self):
+        for constructor in (set, frozenset):
+            t = constructor(self.d.keys())
+            new = marshal.loads(marshal.dumps(t))
+            self.assertEqual(t, new)
+            self.assert_(isinstance(new, constructor))
+            self.assertNotEqual(id(t), id(new))
+            marshal.dump(t, file(test_support.TESTFN, "wb"))
+            marshal.load(file(test_support.TESTFN, "rb"))
+            self.assertEqual(t, new)
+            os.unlink(test_support.TESTFN)
+
 class BugsTestCase(unittest.TestCase):
     def test_bug_5888452(self):
         # Simple-minded check for SF 588452: Debug build crashes
diff --git a/Misc/NEWS b/Misc/NEWS
index 8d00f22..a73ac20 100644
--- a/Misc/NEWS
+++ b/Misc/NEWS
@@ -16,6 +16,8 @@ Core and builtins
 - The peephole optimizer now performs simple constant folding in expressions:
       (2+3) --> (5).
 
+- set and frozenset objects can now be marshalled.  SF #1098985.
+
 
 Extension Modules
 -----------------
diff --git a/Python/marshal.c b/Python/marshal.c
index 0ab0597..2f0a642 100644
--- a/Python/marshal.c
+++ b/Python/marshal.c
@@ -35,6 +35,8 @@
 #define TYPE_CODE	'c'
 #define TYPE_UNICODE	'u'
 #define TYPE_UNKNOWN	'?'
+#define TYPE_SET	'<'
+#define TYPE_FROZENSET  '>'
 
 typedef struct {
 	FILE *fp;
@@ -258,6 +260,38 @@ w_object(PyObject *v, WFILE *p)
 		}
 		w_object((PyObject *)NULL, p);
 	}
+	else if (PyAnySet_Check(v)) {
+		int pos;
+		PyObject *value, *it;
+
+		if (PyObject_TypeCheck(v, &PySet_Type))
+			w_byte(TYPE_SET, p);
+		else
+			w_byte(TYPE_FROZENSET, p);
+		n = PyObject_Size(v);
+		if (n == -1) {
+			p->depth--;
+			p->error = 1;
+			return;
+		}
+		w_long((long)n, p);
+		it = PyObject_GetIter(v);
+		if (it == NULL) {
+			p->depth--;
+			p->error = 1;
+			return;
+		}
+		while ((value = PyIter_Next(it)) != NULL) {
+			w_object(value, p);
+			Py_DECREF(value);
+		}
+		Py_DECREF(it);
+		if (PyErr_Occurred()) {
+			p->depth--;
+			p->error = 1;
+			return;
+		}
+	}
 	else if (PyCode_Check(v)) {
 		PyCodeObject *co = (PyCodeObject *)v;
 		w_byte(TYPE_CODE, p);
@@ -406,7 +440,7 @@ r_object(RFILE *p)
 {
 	/* NULL is a valid return value, it does not necessarily means that
 	   an exception is set. */
-	PyObject *v, *v2;
+	PyObject *v, *v2, *v3;
 	long i, n;
 	int type = r_byte(p);
 
@@ -635,6 +669,37 @@ r_object(RFILE *p)
 		}
 		return v;
 
+	case TYPE_SET:
+	case TYPE_FROZENSET:
+		n = r_long(p);
+		if (n < 0) {
+			PyErr_SetString(PyExc_ValueError, "bad marshal data");
+			return NULL;
+		}
+		v = PyTuple_New((int)n);
+		if (v == NULL)
+			return v;
+		for (i = 0; i < n; i++) {
+			v2 = r_object(p);
+			if ( v2 == NULL ) {
+				if (!PyErr_Occurred())
+					PyErr_SetString(PyExc_TypeError,
+						"NULL object in marshal data");
+				Py_DECREF(v);
+				v = NULL;
+				break;
+			}
+			PyTuple_SET_ITEM(v, (int)i, v2);
+		}
+		if (type == TYPE_SET)
+			v3 = PyObject_CallFunctionObjArgs(
+				(PyObject *)&PySet_Type, v, NULL);
+		else
+			v3 = PyObject_CallFunctionObjArgs(
+				(PyObject *)&PyFrozenSet_Type, v, NULL);
+		Py_DECREF(v);
+		return v3;
+
 	case TYPE_CODE:
 		if (PyEval_GetRestricted()) {
 			PyErr_SetString(PyExc_RuntimeError,