Simplify implementation of choices in GenericReferenceField

This commit is contained in:
Emmanuel Leblond 2015-07-06 10:10:05 +02:00
parent a5fb009b62
commit 34c6790762

View File

@ -1143,32 +1143,27 @@ class GenericReferenceField(BaseField):
def __init__(self, *args, **kwargs): def __init__(self, *args, **kwargs):
choices = kwargs.pop('choices', None) choices = kwargs.pop('choices', None)
super(GenericReferenceField, self).__init__(*args, **kwargs) super(GenericReferenceField, self).__init__(*args, **kwargs)
self._original_choices = choices or [] self.choices = []
self._cooked_choices = None # Keep the choices as a list of allowed Document class names
if choices:
for choice in choices:
if isinstance(choice, basestring):
self.choices.append(choice)
elif isinstance(choice, type) and issubclass(choice, Document):
self.choices.append(choice._class_name)
else:
self.error('Invalid choices provided: must be a list of'
'Document subclasses and/or basestrings')
def _validate_choices(self, value): def _validate_choices(self, value):
if isinstance(value, dict): if isinstance(value, dict):
# If the field has not been dereferenced, it is still a dict # If the field has not been dereferenced, it is still a dict
# of class and DBRef # of class and DBRef
if value.get('_cls') in [c.__name__ for c in self.choices]: value = value.get('_cls')
return elif isinstance(value, Document):
value = value._class_name
super(GenericReferenceField, self)._validate_choices(value) super(GenericReferenceField, self)._validate_choices(value)
@property
def choices(self):
if self._cooked_choices is None:
self._cooked_choices = []
for choice in self._original_choices:
if isinstance(choice, basestring):
choice = get_document(choice)
self._cooked_choices.append(choice)
return self._cooked_choices
@choices.setter
def choices(self, value):
self._original_choices = value
self._cooked_choices = None
def __get__(self, instance, owner): def __get__(self, instance, owner):
if instance is None: if instance is None:
return self return self