Merge pull request #61 from agroszer/master

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,70 @@ 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;
             }
-            if (PyObject_SetItem(*dict, d_key, d_value) < 0)
+        }
+        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);
+            if (items == NULL)
                 return NULL;
+            len = PySequence_Size(items);
+            if (len < 0)
+            {
+                Py_DECREF(items);
+                return NULL;
+            }
+            for ( i=0; i<len; ++i ) {
+                PyObject *item = PySequence_GetItem(items, i);
+                if (item == NULL)
+                {
+                    Py_DECREF(items);
+                    return NULL;
+                }
+                d_key = PyTuple_GetItem(item, 0);
+                if (d_key == NULL)
+                {
+                    Py_DECREF(item);
+                    Py_DECREF(items);
+                    return NULL;
+                }
+                d_value = PyTuple_GetItem(item, 1);
+                if (d_value == NULL)
+                {
+                    Py_DECREF(item);
+                    Py_DECREF(items);
+                    return NULL;
+                }
+
+                if (NATIVE_CHECK_EXACT(d_key)) {
+                    Py_INCREF(d_key);
+                    INTERN_INPLACE(&d_key);
+                    Py_DECREF(d_key);
+                }
+                Py_DECREF(item);
+                if (PyObject_SetItem(*dict, d_key, d_value) < 0)
+                {
+                    Py_DECREF(items);
+                    return NULL;
+                }
+            }
+            Py_DECREF(items);
         }
     }
 

persistent/tests/test_persistence.py

@@ -985,6 +985,23 @@ class _Persistent_Base(object):
             key2 = list(inst2.__dict__.keys())[0]
             self.assertTrue(key1 is key2)
 
+            from persistent._compat import IterableUserDict
+            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 = IterableUserDict({key1: 1})
+            state2 = 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
@@ -998,6 +1015,17 @@ class _Persistent_Base(object):
         inst1.__setstate__({mystr: 2})
         self.assertTrue(mystr in inst1.__dict__)
 
+    def test___setstate___doesnt_fail_on_non_dict(self):
+        class Derived(self._getTargetClass()):
+            pass
+        inst1 = Derived()
+
+        from persistent._compat import IterableUserDict
+        state = IterableUserDict({'foobar': [1, 2]})
+
+        inst1.__setstate__(state)
+        self.assertTrue(hasattr(inst1, 'foobar'))
+
     def test___reduce__(self):
         from persistent._compat import copy_reg
         inst = self._makeOne()