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

This commit is contained in:
Nerl~ 2017-08-11 22:27:25 +03:00 committed by Tim Graham
parent 31f133ea08
commit 97cb3bd16d
2 changed files with 22 additions and 1 deletions

View File

@ -522,11 +522,15 @@ class Model(metaclass=ModelBase):
return hash(self.pk) return hash(self.pk)
def __reduce__(self): def __reduce__(self):
data = self.__dict__ data = self.__getstate__()
data[DJANGO_VERSION_PICKLE_KEY] = get_version() data[DJANGO_VERSION_PICKLE_KEY] = get_version()
class_id = self._meta.app_label, self._meta.object_name class_id = self._meta.app_label, self._meta.object_name
return model_unpickle, (class_id,), data return model_unpickle, (class_id,), data
def __getstate__(self):
"""Hook to allow choosing the attributes to pickle."""
return self.__dict__
def __setstate__(self, state): def __setstate__(self, state):
msg = None msg = None
pickled_version = state.get(DJANGO_VERSION_PICKLE_KEY) pickled_version = state.get(DJANGO_VERSION_PICKLE_KEY)

View File

@ -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() msg = "Pickled model instance's Django version 1.0 does not match the current version %s." % get_version()
with self.assertRaisesMessage(RuntimeWarning, msg): with self.assertRaisesMessage(RuntimeWarning, msg):
pickle.loads(pickle.dumps(p)) 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'))