Fix ``__setstate__`` interning when ``state`` parameter is not a built-in dict

CHANGES.rst

@@ -20,6 +20,8 @@
 
 - Add support for Python 3.6.
 
+- Fix ``__setstate__`` interning when ``state`` parameter is not a built-in dict
+
 
 4.2.2 (2016-11-29)
 ------------------

persistent/cPersistence.c

@@ -588,8 +588,10 @@ pickle___setstate__(PyObject *self, PyObject *state)
     if (state != Py_None)
     {
         PyObject **dict;
+        PyObject *items;
         PyObject *d_key, *d_value;
         Py_ssize_t i;
+        int len;
 
         dict = _PyObject_GetDictPtr(self);
         
@@ -609,17 +611,41 @@ pickle___setstate__(PyObject *self, PyObject *state)
 
         PyDict_Clear(*dict);
 
-        i = 0;
-        while (PyDict_Next(state, &i, &d_key, &d_value)) {
-            /* normally the keys for instance attributes are
-               interned.  we should try to do that here. */
-            if (NATIVE_CHECK_EXACT(d_key)) {
-                Py_INCREF(d_key);
-                INTERN_INPLACE(&d_key);
-                Py_DECREF(d_key);
+        if (PyDict_CheckExact(state))
+        {
+            i = 0;
+            while (PyDict_Next(state, &i, &d_key, &d_value)) {
+                /* normally the keys for instance attributes are
+                   interned.  we should try to do that here. */
+                if (NATIVE_CHECK_EXACT(d_key)) {
+                    Py_INCREF(d_key);
+                    INTERN_INPLACE(&d_key);
+                    Py_DECREF(d_key);
+                }
+                if (PyObject_SetItem(*dict, d_key, d_value) < 0)
+                    return NULL;
+            }
+        }
+        else
+        {
+            /* can happen that not a built-in dict is passed as state
+               fall back to iterating over items, instead of silently
+               failing with PyDict_Next */
+            items = PyMapping_Items(state);
+            len = PySequence_Size(items);
+            for ( i=0; i<len; ++i ) {
+                PyObject *item = PySequence_GetItem(items, i);
+                d_key = PyTuple_GetItem(item, 0);
+                d_value = PyTuple_GetItem(item, 1);
+
+                if (NATIVE_CHECK_EXACT(d_key)) {
+                    Py_INCREF(d_key);
+                    INTERN_INPLACE(&d_key);
+                    Py_DECREF(d_key);
+                }
+                if (PyObject_SetItem(*dict, d_key, d_value) < 0)
+                    return NULL;
             }
-            if (PyObject_SetItem(*dict, d_key, d_value) < 0)
-                return NULL;
         }
     }
 

persistent/tests/test_persistence.py

@@ -985,6 +985,24 @@ class _Persistent_Base(object):
             key2 = list(inst2.__dict__.keys())[0]
             self.assertTrue(key1 is key2)
 
+            import UserDict
+            inst1 = Derived()
+            inst2 = Derived()
+            key1 = 'key'
+            key2 = 'ke'; key2 += 'y'  # construct in a way that won't intern the literal
+            self.assertFalse(key1 is key2)
+            state1 = UserDict.IterableUserDict({key1: 1})
+            state2 = UserDict.IterableUserDict({key2: 2})
+            k1 = list(state1.keys())[0]
+            k2 = list(state2.keys())[0]
+            self.assertFalse(k1 is k2)  # verify
+            inst1.__setstate__(state1)
+            inst2.__setstate__(state2)
+            key1 = list(inst1.__dict__.keys())[0]
+            key2 = list(inst2.__dict__.keys())[0]
+            self.assertTrue(key1 is key2)
+
+
     def test___setstate___doesnt_fail_on_non_string_keys(self):
         class Derived(self._getTargetClass()):
             pass