1 files changed, 373 insertions, 93 deletions

diff --git a/Objects/typeobject.c b/Objects/typeobject.c
index 09f77fa53d..5e951de85e 100644
--- a/Objects/typeobject.c
+++ b/Objects/typeobject.c
@@ -3405,151 +3405,430 @@ import_copyreg(void)
     return cached_copyreg_module;
 }
 
-static PyObject *
-slotnames(PyObject *cls)
+Py_LOCAL(PyObject *)
+_PyType_GetSlotNames(PyTypeObject *cls)
 {
-    PyObject *clsdict;
     PyObject *copyreg;
     PyObject *slotnames;
     _Py_IDENTIFIER(__slotnames__);
     _Py_IDENTIFIER(_slotnames);
 
-    clsdict = ((PyTypeObject *)cls)->tp_dict;
-    slotnames = _PyDict_GetItemId(clsdict, &PyId___slotnames__);
-    if (slotnames != NULL && PyList_Check(slotnames)) {
+    assert(PyType_Check(cls));
+
+    /* Get the slot names from the cache in the class if possible. */
+    slotnames = _PyDict_GetItemIdWithError(cls->tp_dict, &PyId___slotnames__);
+    if (slotnames != NULL) {
+        if (slotnames != Py_None && !PyList_Check(slotnames)) {
+            PyErr_Format(PyExc_TypeError,
+                         "%.200s.__slotnames__ should be a list or None, "
+                         "not %.200s",
+                         cls->tp_name, Py_TYPE(slotnames)->tp_name);
+            return NULL;
+        }
         Py_INCREF(slotnames);
         return slotnames;
     }
+    else {
+        if (PyErr_Occurred()) {
+            return NULL;
+        }
+        /* The class does not have the slot names cached yet. */
+    }
 
     copyreg = import_copyreg();
     if (copyreg == NULL)
         return NULL;
 
-    slotnames = _PyObject_CallMethodId(copyreg, &PyId__slotnames, "O", cls);
+    /* Use _slotnames function from the copyreg module to find the slots
+       by this class and its bases. This function will cache the result
+       in __slotnames__. */
+    slotnames = _PyObject_CallMethodIdObjArgs(copyreg, &PyId__slotnames,
+                                              cls, NULL);
     Py_DECREF(copyreg);
-    if (slotnames != NULL &&
-        slotnames != Py_None &&
-        !PyList_Check(slotnames))
-    {
+    if (slotnames == NULL)
+        return NULL;
+
+    if (slotnames != Py_None && !PyList_Check(slotnames)) {
         PyErr_SetString(PyExc_TypeError,
-            "copyreg._slotnames didn't return a list or None");
+                        "copyreg._slotnames didn't return a list or None");
         Py_DECREF(slotnames);
-        slotnames = NULL;
+        return NULL;
     }
 
     return slotnames;
 }
 
-static PyObject *
-reduce_2(PyObject *obj)
+Py_LOCAL(PyObject *)
+_PyObject_GetState(PyObject *obj)
 {
-    PyObject *cls, *getnewargs;
-    PyObject *args = NULL, *args2 = NULL;
-    PyObject *getstate = NULL, *state = NULL, *names = NULL;
-    PyObject *slots = NULL, *listitems = NULL, *dictitems = NULL;
-    PyObject *copyreg = NULL, *newobj = NULL, *res = NULL;
-    Py_ssize_t i, n;
-    _Py_IDENTIFIER(__getnewargs__);
+    PyObject *state;
+    PyObject *getstate;
     _Py_IDENTIFIER(__getstate__);
-    _Py_IDENTIFIER(__newobj__);
 
-    cls = (PyObject *) Py_TYPE(obj);
+    getstate = _PyObject_GetAttrId(obj, &PyId___getstate__);
+    if (getstate == NULL) {
+        PyObject *slotnames;
 
-    getnewargs = _PyObject_GetAttrId(obj, &PyId___getnewargs__);
-    if (getnewargs != NULL) {
-        args = PyObject_CallObject(getnewargs, NULL);
-        Py_DECREF(getnewargs);
-        if (args != NULL && !PyTuple_Check(args)) {
-            PyErr_Format(PyExc_TypeError,
-                "__getnewargs__ should return a tuple, "
-                "not '%.200s'", Py_TYPE(args)->tp_name);
-            goto end;
+        if (!PyErr_ExceptionMatches(PyExc_AttributeError)) {
+            return NULL;
         }
-    }
-    else {
         PyErr_Clear();
-        args = PyTuple_New(0);
-    }
-    if (args == NULL)
-        goto end;
 
-    getstate = _PyObject_GetAttrId(obj, &PyId___getstate__);
-    if (getstate != NULL) {
-        state = PyObject_CallObject(getstate, NULL);
-        Py_DECREF(getstate);
-        if (state == NULL)
-            goto end;
-    }
-    else {
-        PyObject **dict;
-        PyErr_Clear();
-        dict = _PyObject_GetDictPtr(obj);
-        if (dict && *dict)
-            state = *dict;
-        else
-            state = Py_None;
-        Py_INCREF(state);
-        names = slotnames(cls);
-        if (names == NULL)
-            goto end;
-        if (names != Py_None && PyList_GET_SIZE(names) > 0) {
-            assert(PyList_Check(names));
+        {
+            PyObject **dict;
+            dict = _PyObject_GetDictPtr(obj);
+            /* It is possible that the object's dict is not initialized 
+               yet. In this case, we will return None for the state.
+               We also return None if the dict is empty to make the behavior
+               consistent regardless whether the dict was initialized or not.
+               This make unit testing easier. */
+            if (dict != NULL && *dict != NULL && PyDict_Size(*dict) > 0) {
+                state = *dict;
+            }
+            else {
+                state = Py_None;
+            }
+            Py_INCREF(state);
+        }
+
+        slotnames = _PyType_GetSlotNames(Py_TYPE(obj));
+        if (slotnames == NULL) {
+            Py_DECREF(state);
+            return NULL;
+        }
+
+        assert(slotnames == Py_None || PyList_Check(slotnames));
+        if (slotnames != Py_None && Py_SIZE(slotnames) > 0) {
+            PyObject *slots;
+            Py_ssize_t slotnames_size, i;
+
             slots = PyDict_New();
-            if (slots == NULL)
-                goto end;
-            n = 0;
-            /* Can't pre-compute the list size; the list
-               is stored on the class so accessible to other
-               threads, which may be run by DECREF */
-            for (i = 0; i < PyList_GET_SIZE(names); i++) {
+            if (slots == NULL) {
+                Py_DECREF(slotnames);
+                Py_DECREF(state);
+                return NULL;
+            }
+
+            slotnames_size = Py_SIZE(slotnames);
+            for (i = 0; i < slotnames_size; i++) {
                 PyObject *name, *value;
-                name = PyList_GET_ITEM(names, i);
+
+                name = PyList_GET_ITEM(slotnames, i);
                 value = PyObject_GetAttr(obj, name);
-                if (value == NULL)
+                if (value == NULL) {
+                    if (!PyErr_ExceptionMatches(PyExc_AttributeError)) {
+                        goto error;
+                    }
+                    /* It is not an error if the attribute is not present. */
                     PyErr_Clear();
+                }
                 else {
-                    int err = PyDict_SetItem(slots, name,
-                                             value);
+                    int err = PyDict_SetItem(slots, name, value);
                     Py_DECREF(value);
-                    if (err)
-                        goto end;
-                    n++;
+                    if (err) {
+                        goto error;
+                    }
+                }
+
+                /* The list is stored on the class so it may mutates while we
+                   iterate over it */
+                if (slotnames_size != Py_SIZE(slotnames)) {
+                    PyErr_Format(PyExc_RuntimeError,
+                                 "__slotsname__ changed size during iteration");
+                    goto error;
+                }
+
+                /* We handle errors within the loop here. */
+                if (0) {
+                  error:
+                    Py_DECREF(slotnames);
+                    Py_DECREF(slots);
+                    Py_DECREF(state);
+                    return NULL;
                 }
             }
-            if (n) {
-                state = Py_BuildValue("(NO)", state, slots);
-                if (state == NULL)
-                    goto end;
+
+            /* If we found some slot attributes, pack them in a tuple along
+               the orginal attribute dictionary. */
+            if (PyDict_Size(slots) > 0) {
+                PyObject *state2;
+
+                state2 = PyTuple_Pack(2, state, slots);
+                Py_DECREF(state);
+                if (state2 == NULL) {
+                    Py_DECREF(slotnames);
+                    Py_DECREF(slots);
+                    return NULL;
+                }
+                state = state2;
             }
+            Py_DECREF(slots);
+        }
+        Py_DECREF(slotnames);
+    }
+    else { /* getstate != NULL */
+        state = PyObject_CallObject(getstate, NULL);
+        Py_DECREF(getstate);
+        if (state == NULL)
+            return NULL;
+    }
+
+    return state;
+}
+
+Py_LOCAL(int)
+_PyObject_GetNewArguments(PyObject *obj, PyObject **args, PyObject **kwargs)
+{
+    PyObject *getnewargs, *getnewargs_ex;
+    _Py_IDENTIFIER(__getnewargs_ex__);
+    _Py_IDENTIFIER(__getnewargs__);
+
+    if (args == NULL || kwargs == NULL) {
+        PyErr_BadInternalCall();
+        return -1;
+    }
+
+    /* We first attempt to fetch the arguments for __new__ by calling
+       __getnewargs_ex__ on the object. */
+    getnewargs_ex = _PyObject_GetAttrId(obj, &PyId___getnewargs_ex__);
+    if (getnewargs_ex != NULL) {
+        PyObject *newargs = PyObject_CallObject(getnewargs_ex, NULL);
+        Py_DECREF(getnewargs_ex);
+        if (newargs == NULL) {
+            return -1;
+        }
+        if (!PyTuple_Check(newargs)) {
+            PyErr_Format(PyExc_TypeError,
+                         "__getnewargs_ex__ should return a tuple, "
+                         "not '%.200s'", Py_TYPE(newargs)->tp_name);
+            Py_DECREF(newargs);
+            return -1;
+        }
+        if (Py_SIZE(newargs) != 2) {
+            PyErr_Format(PyExc_ValueError,
+                         "__getnewargs_ex__ should return a tuple of "
+                         "length 2, not %zd", Py_SIZE(newargs));
+            Py_DECREF(newargs);
+            return -1;
+        }
+        *args = PyTuple_GET_ITEM(newargs, 0);
+        Py_INCREF(*args);
+        *kwargs = PyTuple_GET_ITEM(newargs, 1);
+        Py_INCREF(*kwargs);
+        Py_DECREF(newargs);
+
+        /* XXX We should perhaps allow None to be passed here. */
+        if (!PyTuple_Check(*args)) {
+            PyErr_Format(PyExc_TypeError,
+                         "first item of the tuple returned by "
+                         "__getnewargs_ex__ must be a tuple, not '%.200s'",
+                         Py_TYPE(*args)->tp_name);
+            Py_CLEAR(*args);
+            Py_CLEAR(*kwargs);
+            return -1;
+        }
+        if (!PyDict_Check(*kwargs)) {
+            PyErr_Format(PyExc_TypeError,
+                         "second item of the tuple returned by "
+                         "__getnewargs_ex__ must be a dict, not '%.200s'",
+                         Py_TYPE(*kwargs)->tp_name);
+            Py_CLEAR(*args);
+            Py_CLEAR(*kwargs);
+            return -1;
+        }
+        return 0;
+    } else {
+        if (!PyErr_ExceptionMatches(PyExc_AttributeError)) {
+            return -1;
+        }
+        PyErr_Clear();
+    }
+
+    /* The object does not have __getnewargs_ex__ so we fallback on using
+       __getnewargs__ instead. */
+    getnewargs = _PyObject_GetAttrId(obj, &PyId___getnewargs__);
+    if (getnewargs != NULL) {
+        *args = PyObject_CallObject(getnewargs, NULL);
+        Py_DECREF(getnewargs);
+        if (*args == NULL) {
+            return -1;
+        }
+        if (!PyTuple_Check(*args)) {
+            PyErr_Format(PyExc_TypeError,
+                         "__getnewargs__ should return a tuple, "
+                         "not '%.200s'", Py_TYPE(*args)->tp_name);
+            Py_CLEAR(*args);
+            return -1;
+        }
+        *kwargs = NULL;
+        return 0;
+    } else {
+        if (!PyErr_ExceptionMatches(PyExc_AttributeError)) {
+            return -1;
         }
+        PyErr_Clear();
+    }
+
+    /* The object does not have __getnewargs_ex__ and __getnewargs__. This may
+       means __new__ does not takes any arguments on this object, or that the
+       object does not implement the reduce protocol for pickling or
+       copying. */
+    *args = NULL;
+    *kwargs = NULL;
+    return 0;
+}
+
+Py_LOCAL(int)
+_PyObject_GetItemsIter(PyObject *obj, PyObject **listitems,
+                       PyObject **dictitems)
+{
+    if (listitems == NULL || dictitems == NULL) {
+        PyErr_BadInternalCall();
+        return -1;
     }
 
     if (!PyList_Check(obj)) {
-        listitems = Py_None;
-        Py_INCREF(listitems);
+        *listitems = Py_None;
+        Py_INCREF(*listitems);
     }
     else {
-        listitems = PyObject_GetIter(obj);
+        *listitems = PyObject_GetIter(obj);
         if (listitems == NULL)
-            goto end;
+            return -1;
     }
 
     if (!PyDict_Check(obj)) {
-        dictitems = Py_None;
-        Py_INCREF(dictitems);
+        *dictitems = Py_None;
+        Py_INCREF(*dictitems);
     }
     else {
+        PyObject *items;
         _Py_IDENTIFIER(items);
-        PyObject *items = _PyObject_CallMethodId(obj, &PyId_items, "");
-        if (items == NULL)
-            goto end;
-        dictitems = PyObject_GetIter(items);
+
+        items = _PyObject_CallMethodIdObjArgs(obj, &PyId_items, NULL);
+        if (items == NULL) {
+            Py_CLEAR(*listitems);
+            return -1;
+        }
+        *dictitems = PyObject_GetIter(items);
         Py_DECREF(items);
-        if (dictitems == NULL)
-            goto end;
+        if (*dictitems == NULL) {
+            Py_CLEAR(*listitems);
+            return -1;
+        }
+    }
+
+    assert(*listitems != NULL && *dictitems != NULL);
+
+    return 0;
+}
+
+static PyObject *
+reduce_4(PyObject *obj)
+{
+    PyObject *args = NULL, *kwargs = NULL;
+    PyObject *copyreg;
+    PyObject *newobj, *newargs, *state, *listitems, *dictitems;
+    PyObject *result;
+    _Py_IDENTIFIER(__newobj_ex__);
+
+    if (_PyObject_GetNewArguments(obj, &args, &kwargs) < 0) {
+        return NULL;
+    }
+    if (args == NULL) {
+        args = PyTuple_New(0);
+        if (args == NULL)
+            return NULL;
+    }
+    if (kwargs == NULL) {
+        kwargs = PyDict_New();
+        if (kwargs == NULL)
+            return NULL;
     }
 
     copyreg = import_copyreg();
+    if (copyreg == NULL) {
+        Py_DECREF(args);
+        Py_DECREF(kwargs);
+        return NULL;
+    }
+    newobj = _PyObject_GetAttrId(copyreg, &PyId___newobj_ex__);
+    Py_DECREF(copyreg);
+    if (newobj == NULL) {
+        Py_DECREF(args);
+        Py_DECREF(kwargs);
+        return NULL;
+    }
+    newargs = PyTuple_Pack(3, Py_TYPE(obj), args, kwargs);
+    Py_DECREF(args);
+    Py_DECREF(kwargs);
+    if (newargs == NULL) {
+        Py_DECREF(newobj);
+        return NULL;
+    }
+    state = _PyObject_GetState(obj);
+    if (state == NULL) {
+        Py_DECREF(newobj);
+        Py_DECREF(newargs);
+        return NULL;
+    }
+    if (_PyObject_GetItemsIter(obj, &listitems, &dictitems) < 0) {
+        Py_DECREF(newobj);
+        Py_DECREF(newargs);
+        Py_DECREF(state);
+        return NULL;
+    }
+
+    result = PyTuple_Pack(5, newobj, newargs, state, listitems, dictitems);
+    Py_DECREF(newobj);
+    Py_DECREF(newargs);
+    Py_DECREF(state);
+    Py_DECREF(listitems);
+    Py_DECREF(dictitems);
+    return result;    
+}
+
+static PyObject *
+reduce_2(PyObject *obj)
+{
+    PyObject *cls;
+    PyObject *args = NULL, *args2 = NULL, *kwargs = NULL;
+    PyObject *state = NULL, *listitems = NULL, *dictitems = NULL;
+    PyObject *copyreg = NULL, *newobj = NULL, *res = NULL;
+    Py_ssize_t i, n;
+    _Py_IDENTIFIER(__newobj__);
+
+    if (_PyObject_GetNewArguments(obj, &args, &kwargs) < 0) {
+        return NULL;
+    }
+    if (args == NULL) {
+        assert(kwargs == NULL);
+        args = PyTuple_New(0);
+        if (args == NULL) {
+            return NULL;
+        }
+    }
+    else if (kwargs != NULL) {
+        if (PyDict_Size(kwargs) > 0) {
+            PyErr_SetString(PyExc_ValueError, 
+                            "must use protocol 4 or greater to copy this "
+                            "object; since __getnewargs_ex__ returned "
+                            "keyword arguments.");
+            Py_DECREF(args);
+            Py_DECREF(kwargs);
+            return NULL;
+        }
+        Py_CLEAR(kwargs);
+    }
+
+    state = _PyObject_GetState(obj);
+    if (state == NULL)
+        goto end;
+
+    if (_PyObject_GetItemsIter(obj, &listitems, &dictitems) < 0)
+        goto end;
+
+    copyreg = import_copyreg();
     if (copyreg == NULL)
         goto end;
     newobj = _PyObject_GetAttrId(copyreg, &PyId___newobj__);
@@ -3560,6 +3839,7 @@ reduce_2(PyObject *obj)
     args2 = PyTuple_New(n+1);
     if (args2 == NULL)
         goto end;
+    cls = (PyObject *) Py_TYPE(obj);
     Py_INCREF(cls);
     PyTuple_SET_ITEM(args2, 0, cls);
     for (i = 0; i < n; i++) {
@@ -3573,9 +3853,7 @@ reduce_2(PyObject *obj)
   end:
     Py_XDECREF(args);
     Py_XDECREF(args2);
-    Py_XDECREF(slots);
     Py_XDECREF(state);
-    Py_XDECREF(names);
     Py_XDECREF(listitems);
     Py_XDECREF(dictitems);
     Py_XDECREF(copyreg);
@@ -3603,7 +3881,9 @@ _common_reduce(PyObject *self, int proto)
 {
     PyObject *copyreg, *res;
 
-    if (proto >= 2)
+    if (proto >= 4)
+        return reduce_4(self);
+    else if (proto >= 2)
         return reduce_2(self);
 
     copyreg = import_copyreg();