From 97cb3bd16d8e50d96eb184f8d67dd1723f776d2a Mon Sep 17 00:00:00 2001 From: Nerl~ Date: Fri, 11 Aug 2017 22:27:25 +0300 Subject: [PATCH] Fixed #28456 -- Allowed customizing Model pickling by overriding __getstate__(). --- django/db/models/base.py | 6 +++++- tests/model_regress/test_pickle.py | 17 +++++++++++++++++ 2 files changed, 22 insertions(+), 1 deletion(-) diff --git a/django/db/models/base.py b/django/db/models/base.py index dda2248308c..3cff2d2fc42 100644 --- a/django/db/models/base.py +++ b/django/db/models/base.py @@ -522,11 +522,15 @@ class Model(metaclass=ModelBase): 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 index a3df362e324..5fbe0a3cbbb 100644 --- a/tests/model_regress/test_pickle.py +++ b/tests/model_regress/test_pickle.py @@ -43,3 +43,20 @@ class ModelPickleTestCase(TestCase): 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'))