Fixes pickling issue with choice fields

Removes the dynamic __get_field_display partials before pickling
This commit is contained in:
Ross Lawley 2011-06-16 12:50:45 +01:00
parent 22a7ee5885
commit cae3f3eeff
2 changed files with 51 additions and 27 deletions

View File

@ -614,9 +614,6 @@ class BaseDocument(object):
self._data = {}
# Assign default values to instance
for attr_name, field in self._fields.items():
if field.choices: # dynamically adds a way to get the display value for a field with choices
setattr(self, 'get_%s_display' % attr_name, partial(self._get_FIELD_display, field=field))
value = getattr(self, attr_name, None)
setattr(self, attr_name, value)
@ -628,9 +625,29 @@ class BaseDocument(object):
except AttributeError:
pass
# Set any get_fieldname_display methods
self.__set_field_display()
signals.post_init.send(self.__class__, document=self)
def _get_FIELD_display(self, field):
def __getstate__(self):
self_dict = self.__dict__
removals = ["get_%s_display" % k for k,v in self._fields.items() if v.choices]
for k in removals:
if hasattr(self, k):
delattr(self, k)
return self.__dict__
def __setstate__(self, __dict__):
self.__dict__ = __dict__
self.__set_field_display()
def __set_field_display(self):
for attr_name, field in self._fields.items():
if field.choices: # dynamically adds a way to get the display value for a field with choices
setattr(self, 'get_%s_display' % attr_name, partial(self.__get_field_display, field=field))
def __get_field_display(self, field):
"""Returns the display value for a choice field"""
value = getattr(self, field.name)
return dict(field.choices).get(value, value)
@ -865,42 +882,46 @@ class BaseList(list):
super(BaseList, self).__init__(list_items)
def __setitem__(self, *args, **kwargs):
if hasattr(self, 'instance') and hasattr(self, 'name'):
self.instance._mark_as_changed(self.name)
self._mark_as_changed()
super(BaseList, self).__setitem__(*args, **kwargs)
def __delitem__(self, *args, **kwargs):
self.instance._mark_as_changed(self.name)
self._mark_as_changed()
super(BaseList, self).__delitem__(*args, **kwargs)
def append(self, *args, **kwargs):
self.instance._mark_as_changed(self.name)
self._mark_as_changed()
return super(BaseList, self).append(*args, **kwargs)
def extend(self, *args, **kwargs):
self.instance._mark_as_changed(self.name)
self._mark_as_changed()
return super(BaseList, self).extend(*args, **kwargs)
def insert(self, *args, **kwargs):
self.instance._mark_as_changed(self.name)
self._mark_as_changed()
return super(BaseList, self).insert(*args, **kwargs)
def pop(self, *args, **kwargs):
self.instance._mark_as_changed(self.name)
self._mark_as_changed()
return super(BaseList, self).pop(*args, **kwargs)
def remove(self, *args, **kwargs):
self.instance._mark_as_changed(self.name)
self._mark_as_changed()
return super(BaseList, self).remove(*args, **kwargs)
def reverse(self, *args, **kwargs):
self.instance._mark_as_changed(self.name)
self._mark_as_changed()
return super(BaseList, self).reverse(*args, **kwargs)
def sort(self, *args, **kwargs):
self.instance._mark_as_changed(self.name)
self._mark_as_changed()
return super(BaseList, self).sort(*args, **kwargs)
def _mark_as_changed(self):
"""Marks a list as changed if has an instance and a name"""
if hasattr(self, 'instance') and hasattr(self, 'name'):
self.instance._mark_as_changed(self.name)
class BaseDict(dict):
"""A special dict so we can watch any changes
@ -912,39 +933,42 @@ class BaseDict(dict):
super(BaseDict, self).__init__(dict_items)
def __setitem__(self, *args, **kwargs):
if hasattr(self, 'instance') and hasattr(self, 'name'):
self.instance._mark_as_changed(self.name)
self._mark_as_changed()
super(BaseDict, self).__setitem__(*args, **kwargs)
def __setattr__(self, *args, **kwargs):
if hasattr(self, 'instance') and hasattr(self, 'name'):
self.instance._mark_as_changed(self.name)
self._mark_as_changed()
super(BaseDict, self).__setattr__(*args, **kwargs)
def __delete__(self, *args, **kwargs):
self.instance._mark_as_changed(self.name)
self._mark_as_changed()
super(BaseDict, self).__delete__(*args, **kwargs)
def __delitem__(self, *args, **kwargs):
self.instance._mark_as_changed(self.name)
self._mark_as_changed()
super(BaseDict, self).__delitem__(*args, **kwargs)
def __delattr__(self, *args, **kwargs):
self.instance._mark_as_changed(self.name)
self._mark_as_changed()
super(BaseDict, self).__delattr__(*args, **kwargs)
def clear(self, *args, **kwargs):
self.instance._mark_as_changed(self.name)
self._mark_as_changed()
super(BaseDict, self).clear(*args, **kwargs)
def pop(self, *args, **kwargs):
self.instance._mark_as_changed(self.name)
self._mark_as_changed()
super(BaseDict, self).clear(*args, **kwargs)
def popitem(self, *args, **kwargs):
self.instance._mark_as_changed(self.name)
self._mark_as_changed()
super(BaseDict, self).clear(*args, **kwargs)
def _mark_as_changed(self):
"""Marks a dict as changed if has an instance and a name"""
if hasattr(self, 'instance') and hasattr(self, 'name'):
self.instance._mark_as_changed(self.name)
if sys.version_info < (2, 5):
# Prior to Python 2.5, Exception was an old-style class
import types

View File

@ -15,7 +15,7 @@ class PickleEmbedded(EmbeddedDocument):
class PickleTest(Document):
number = IntField()
string = StringField()
string = StringField(choices=(('One', '1'), ('Two', '2')))
embedded = EmbeddedDocumentField(PickleEmbedded)
lists = ListField(StringField())
@ -1516,7 +1516,7 @@ class DocumentTest(unittest.TestCase):
def test_picklable(self):
pickle_doc = PickleTest(number=1, string="OH HAI", lists=['1', '2'])
pickle_doc = PickleTest(number=1, string="One", lists=['1', '2'])
pickle_doc.embedded = PickleEmbedded()
pickle_doc.save()
@ -1525,7 +1525,7 @@ class DocumentTest(unittest.TestCase):
self.assertEquals(resurrected, pickle_doc)
resurrected.string = "Working"
resurrected.string = "Two"
resurrected.save()
pickle_doc.reload()