Merge pull request #16 from davisagli/davisagli-intern-state-keys

Intern keys in `__setstate__`

CHANGES.rst

@@ -4,6 +4,9 @@
 4.0.9 (unreleased)
 ------------------
 
+- Intern keys of object state in `__setstate__` to reduce memory usage
+  when unpickling multiple objects with the same attributes.
+
 - Add support for PyPy3.
 
 - 100% branch coverage.

persistent/_compat.h

@@ -23,6 +23,8 @@
 
 #ifdef PY3K
 #define INTERN PyUnicode_InternFromString
+#define INTERN_INPLACE PyUnicode_InternInPlace
+#define NATIVE_CHECK_EXACT PyUnicode_CheckExact
 #define NATIVE_FROM_STRING_AND_SIZE PyUnicode_FromStringAndSize
 
 #define Py_TPFLAGS_HAVE_RICHCOMPARE 0
@@ -34,6 +36,8 @@
 
 #else
 #define INTERN PyString_InternFromString
+#define INTERN_INPLACE PyString_InternInPlace
+#define NATIVE_CHECK_EXACT PyString_CheckExact
 #define NATIVE_FROM_STRING_AND_SIZE PyString_FromStringAndSize
 
 #define INT_FROM_LONG(x) PyInt_FromLong(x)

persistent/_compat.py

@@ -18,6 +18,7 @@
     import copyreg as copy_reg
     from collections import UserDict as IterableUserDict
     from collections import UserList
+    from sys import intern
 
     def _u(s):
         return s
@@ -52,3 +53,5 @@ def _native(s):
 
     PYTHON3 = False
     PYTHON2 = True
+
+    intern = intern

persistent/cPersistence.c

@@ -534,6 +534,8 @@ pickle___setstate__(PyObject *self, PyObject *state)
     if (state != Py_None)
     {
         PyObject **dict;
+        PyObject *d_key, *d_value;
+        Py_ssize_t i;
 
         dict = _PyObject_GetDictPtr(self);
 
@@ -552,8 +554,19 @@ pickle___setstate__(PyObject *self, PyObject *state)
         }
 
         PyDict_Clear(*dict);
-        if (PyDict_Update(*dict, state) < 0)
-            return NULL;
+
+        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 (slots && pickle_setattrs_from_dict(self, slots) < 0)

persistent/persistence.py

@@ -25,6 +25,7 @@
 from persistent.timestamp import TimeStamp
 from persistent.timestamp import _ZERO
 from persistent._compat import copy_reg
+from persistent._compat import intern
 
 _INITIAL_SERIAL = _ZERO
 
@@ -313,7 +314,8 @@ def __setstate__(self, state):
             if idict is None:
                 raise TypeError('No instance dict')
             idict.clear()
-            idict.update(inst_dict)
+            for k, v in inst_dict.items():
+                idict[intern(k)] = v
         slotnames = self._slotnames()
         if slotnames:
             for k, v in slots.items():

persistent/tests/test_persistence.py

@@ -863,6 +863,20 @@ class Derived(Base):
         self.assertEqual(inst.baz, 'bam')
         self.assertEqual(inst.qux, 'spam')
 
+    def test___setstate___interns_dict_keys(self):
+        class Derived(self._getTargetClass()):
+            pass
+        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)
+        inst1.__setstate__({key1: 1})
+        inst2.__setstate__({key2: 2})
+        key1 = list(inst1.__dict__.keys())[0]
+        key2 = list(inst2.__dict__.keys())[0]
+        self.assertTrue(key1 is key2)
+
     def test___reduce__(self):
         from persistent._compat import copy_reg
         inst = self._makeOne()