Fixed #28456 -- Allowed customizing Model pickling by overriding __getstate__().

Nerlin · timgraham · commit 97cb3bd16d8e · 2017-08-11T15:27:25.000-04:00
diff --git a/django/db/models/base.py b/django/db/models/base.py
@@ -522,11 +522,15 @@ def __hash__(self):
         return hash(self.pk)
 
     def __reduce__(self):
-        data = self.__dict__
+        data = self.__getstate__()
         data[DJANGO_VERSION_PICKLE_KEY] = get_version()
         class_id = self._meta.app_label, self._meta.object_name
         return model_unpickle, (class_id,), data
 
+    def __getstate__(self):
+        """Hook to allow choosing the attributes to pickle."""
+        return self.__dict__
+
     def __setstate__(self, state):
         msg = None
         pickled_version = state.get(DJANGO_VERSION_PICKLE_KEY)
diff --git a/tests/model_regress/test_pickle.py b/tests/model_regress/test_pickle.py
@@ -43,3 +43,20 @@ def __reduce__(self):
         msg = "Pickled model instance's Django version 1.0 does not match the current version %s." % get_version()
         with self.assertRaisesMessage(RuntimeWarning, msg):
             pickle.loads(pickle.dumps(p))
+
+    def test_with_getstate(self):
+        """
+        A model may override __getstate__() to choose the attributes to pickle.
+        """
+        class PickledModel(models.Model):
+            def __getstate__(self):
+                state = super().__getstate__().copy()
+                del state['dont_pickle']
+                return state
+
+        m = PickledModel()
+        m.dont_pickle = 1
+        dumped = pickle.dumps(m)
+        self.assertEqual(m.dont_pickle, 1)
+        reloaded = pickle.loads(dumped)
+        self.assertFalse(hasattr(reloaded, 'dont_pickle'))
