SF 1062353: set pickling problems

Support automatic pickling of dictionaries in instance of set subclasses.

2 files changed, 18 insertions, 5 deletions

diff --git a/Lib/test/test_set.py b/Lib/test/test_set.py
index aab0c57..900d8e4 100644
--- a/Lib/test/test_set.py
+++ b/Lib/test/test_set.py
@@ -175,9 +175,15 @@ class TestJointOps(unittest.TestCase):
         self.failIf(set('cbs').issuperset('a'))
 
     def test_pickling(self):
-        p = pickle.dumps(self.s)
-        dup = pickle.loads(p)
-        self.assertEqual(self.s, dup, "%s != %s" % (self.s, dup))
+        for i in (0, 1, 2):
+            p = pickle.dumps(self.s, i)
+            dup = pickle.loads(p)
+            self.assertEqual(self.s, dup, "%s != %s" % (self.s, dup))
+            if type(self.s) not in (set, frozenset):
+                self.s.x = 10
+                p = pickle.dumps(self.s)
+                dup = pickle.loads(p)
+                self.assertEqual(self.s.x, dup.x)
 
     def test_deepcopy(self):
         class Tracer:
diff --git a/Objects/setobject.c b/Objects/setobject.c
index d57217c..8ef671e 100644
--- a/Objects/setobject.c
+++ b/Objects/setobject.c
@@ -844,7 +844,7 @@ PyDoc_STRVAR(pop_doc, "Remove and return an arbitrary set element.");
 static PyObject *
 set_reduce(PySetObject *so)
 {
-	PyObject *keys=NULL, *args=NULL, *result=NULL;
+	PyObject *keys=NULL, *args=NULL, *result=NULL, *dict=NULL;
 
 	keys = PyDict_Keys(so->data);
 	if (keys == NULL)
@@ -852,10 +852,17 @@ set_reduce(PySetObject *so)
 	args = PyTuple_Pack(1, keys);
 	if (args == NULL)
 		goto done;
-	result = PyTuple_Pack(2, so->ob_type, args);
+	dict = PyObject_GetAttrString((PyObject *)so, "__dict__");
+	if (dict == NULL) {
+		PyErr_Clear();
+		dict = Py_None;
+		Py_INCREF(dict);
+	}
+	result = PyTuple_Pack(3, so->ob_type, args, dict);
 done:
 	Py_XDECREF(args);
 	Py_XDECREF(keys);
+	Py_XDECREF(dict);
 	return result;
 }