# HG changeset patch
# Parent c56a74b0ebf0dbd77b6ef5785f9a7a17e93dbe2f
diff -r c56a74b0ebf0 -r c0b803209dc8 Objects/setobject.c
--- a/Objects/setobject.c	Fri Jul 27 00:00:51 2012 +0300
+++ b/Objects/setobject.c	Fri Jul 27 02:49:05 2012 +0300
@@ -1991,31 +1991,6 @@
 \n\
 If the element is not a member, do nothing.");
 
-static PyObject *
-set_reduce(PySetObject *so)
-{
-    PyObject *keys=NULL, *args=NULL, *result=NULL, *dict=NULL;
-    _Py_IDENTIFIER(__dict__);
-
-    keys = PySequence_List((PyObject *)so);
-    if (keys == NULL)
-        goto done;
-    args = PyTuple_Pack(1, keys);
-    if (args == NULL)
-        goto done;
-    dict = _PyObject_GetAttrId((PyObject *)so, &PyId___dict__);
-    if (dict == NULL) {
-        PyErr_Clear();
-        dict = Py_None;
-        Py_INCREF(dict);
-    }
-    result = PyTuple_Pack(3, Py_TYPE(so), args, dict);
-done:
-    Py_XDECREF(args);
-    Py_XDECREF(keys);
-    Py_XDECREF(dict);
-    return result;
-}
 
 static PyObject *
 set_sizeof(PySetObject *so)
@@ -2047,6 +2022,124 @@
     return set_update_internal(self, iterable);
 }
 
+//returns either a list or a tuple of the form (list, dict)
+//list is the list of elements currently in the set and dict is the value
+//of so.__dict__, if there is one
+static PyObject *
+set_getstate(PySetObject *so)
+{
+    PyObject *keys, *dict, *state = NULL;
+    _Py_IDENTIFIER(__dict__);
+
+    keys = PySequence_List((PyObject *)so);
+    if (keys == NULL) return NULL;
+
+    dict = _PyObject_GetAttrId((PyObject *)so, &PyId___dict__);
+    if (dict == NULL || PyDict_Size(dict) == 0) {
+        //no __dict__, the state is a list
+        PyErr_Clear();
+        state = keys;
+        Py_XDECREF(dict);
+    }
+    else {
+        //got a __dict__, the state is a 2-tuple
+        state = PyTuple_Pack(2, keys, dict);
+        Py_DECREF(keys);
+        Py_DECREF(dict);
+    }
+    return state;
+}
+
+static PyObject *
+set_reduce(PySetObject *so)
+{
+    PyObject *empty_tuple = PyTuple_New(0);
+    PyObject *state;
+    PyObject *ret;
+
+    empty_tuple = PyTuple_New(0);
+    if (!empty_tuple) return NULL;
+
+    state = set_getstate(so);
+    if (!state) return NULL;
+
+    ret = PyTuple_Pack(3, Py_TYPE(so), empty_tuple, state);
+    Py_DECREF(empty_tuple);
+    Py_DECREF(state);
+
+    return ret;
+}
+
+PyDoc_STRVAR(setstate_doc, "Update a set from an unpickled state");
+static PyObject *
+set_setstate(PySetObject *so, PyObject *state)
+{
+    PyObject *keys, *dict = NULL;
+    PyObject *it, *key;
+    int is_list;
+    _Py_IDENTIFIER(__dict__);
+    
+    is_list = PyList_Check(state);
+    if (!is_list && !PyTuple_Check(state)) {
+        PyErr_Format(PyExc_TypeError, 
+                "set.__setstate__ expected a tuple or list, got a '%.200s'",
+                state->ob_type->tp_name);
+        return NULL;
+    }
+
+    if (is_list) keys = state;
+    else {
+        if (!PyArg_UnpackTuple(state, 
+                    "set.__setstate__'s tuple", 2, 2, &keys, &dict))
+            return NULL;
+
+        if (!PyList_Check(keys)) {
+            PyErr_Format(PyExc_TypeError,
+                    "set.__setstate__ expected the tuple's first parameter "
+                    "to be a list, got a '%.200s'",
+                    keys->ob_type->tp_name);
+            return NULL;
+        }
+
+        if (!PyDict_Check(dict)) {
+            PyErr_Format(PyExc_TypeError,
+                    "set.__setstate__ expected the tuple's second parameter "
+                    "to be a dict, got a '%.200s'",
+                    dict->ob_type->tp_name);
+            return NULL;
+        }
+
+        if (!_PyObject_HasAttrId((PyObject *)so, &PyId___dict__)) {
+            PyErr_SetString(PyExc_AttributeError, 
+                    "set.__setstate__ was provided with a dict "
+                    "when the underlying object does not have a __dict__ "
+                    "attribute.");
+            return NULL;
+        }
+    }
+
+    set_clear_internal(so);
+
+    it = PyObject_GetIter(keys);
+    if (it == NULL) return NULL;
+
+    while ((key = PyIter_Next(it)) != NULL) {
+        if (set_add_key(so, key) == -1) {
+            Py_DECREF(key);
+            Py_DECREF(it);
+            return NULL;
+        }
+        Py_DECREF(key);
+    }
+    Py_DECREF(it);
+    if (PyErr_Occurred()) return NULL;
+
+    if (dict && _PyObject_SetAttrId((PyObject *)so, &PyId___dict__, dict) < 0)
+        return NULL;
+
+    Py_RETURN_NONE;
+}
+
 static PySequenceMethods set_as_sequence = {
     set_len,                            /* sq_length */
     0,                                  /* sq_concat */
@@ -2094,6 +2187,10 @@
      issuperset_doc},
     {"pop",             (PyCFunction)set_pop,           METH_NOARGS,
      pop_doc},
+    {"__getstate__",    (PyCFunction)set_getstate,      METH_NOARGS,
+     reduce_doc},
+    {"__setstate__",    (PyCFunction)set_setstate,      METH_O,
+     setstate_doc},
     {"__reduce__",      (PyCFunction)set_reduce,        METH_NOARGS,
      reduce_doc},
     {"remove",          (PyCFunction)set_remove,        METH_O,
@@ -2199,6 +2296,31 @@
 
 /* frozenset object ********************************************************/
 
+static PyObject *
+frozenset_reduce(PySetObject *so)
+{
+    PyObject *keys=NULL, *args=NULL, *result=NULL, *dict=NULL;
+    _Py_IDENTIFIER(__dict__);
+
+    keys = PySequence_List((PyObject *)so);
+    if (keys == NULL)
+        goto done;
+    args = PyTuple_Pack(1, keys);
+    if (args == NULL)
+        goto done;
+    dict = _PyObject_GetAttrId((PyObject *)so, &PyId___dict__);
+    if (dict == NULL) {
+        PyErr_Clear();
+        dict = Py_None;
+        Py_INCREF(dict);
+    }
+    result = PyTuple_Pack(3, Py_TYPE(so), args, dict);
+done:
+    Py_XDECREF(args);
+    Py_XDECREF(keys);
+    Py_XDECREF(dict);
+    return result;
+}
 
 static PyMethodDef frozenset_methods[] = {
     {"__contains__",(PyCFunction)set_direct_contains,           METH_O | METH_COEXIST,
@@ -2215,7 +2337,7 @@
      issubset_doc},
     {"issuperset",      (PyCFunction)set_issuperset,    METH_O,
      issuperset_doc},
-    {"__reduce__",      (PyCFunction)set_reduce,        METH_NOARGS,
+    {"__reduce__",      (PyCFunction)frozenset_reduce,  METH_NOARGS,
      reduce_doc},
     {"__sizeof__",      (PyCFunction)set_sizeof,        METH_NOARGS,
      sizeof_doc},