SF 1098985: set objects cannot be marshalled

diff --git a/Lib/test/test_marshal.py b/Lib/test/test_marshal.py
index 5c1a3f3c6d27d333f90dfb8dcc0ce8b690e1a68e..9901a3c2d4eac366cc4a2397950eba18954ceeaf 100644 (file)
--- 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 8d00f22a6436c3e6e5946a4cfd1766b8e3e22be9..a73ac206f93a417dae4aca59b4d9cfd451cfd8ae 100644 (file)
--- 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 0ab0597d321bbb59b62bdc1e5b81a1bc93f82ce9..2f0a6427a6ae933420ee40f9b88b6505fee2173a 100644 (file)
--- 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,