diff options
| -rw-r--r-- | Lib/test/test_marshal.py | 12 | ||||
| -rw-r--r-- | Misc/NEWS | 2 | ||||
| -rw-r--r-- | Python/marshal.c | 67 |
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 @@ -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, |