diff --git a/mongoengine/base/fields.py b/mongoengine/base/fields.py index 9ba9dc9a..6623d1f2 100644 --- a/mongoengine/base/fields.py +++ b/mongoengine/base/fields.py @@ -189,14 +189,18 @@ class BaseField(object): pass def _validate_choices(self, value): + """Validate that value is a valid choice for this field.""" Document = _import_class('Document') EmbeddedDocument = _import_class('EmbeddedDocument') + # Field choices can be given as an iterable (e.g. tuple/list/set) of + # values or an iterable of value-label pairs, e.g. ('XS', 'Extra Small'). + # It the latter case, extract just the values for comparison. choice_list = self.choices - if isinstance(choice_list[0], (list, tuple)): - choice_list = [k for k, _ in choice_list] + if isinstance(next(iter(choice_list)), (list, tuple)): + choice_list = [val for val, label in choice_list] - # Choices which are other types of Documents + # Validate Document/EmbeddedDocument choices if isinstance(value, (Document, EmbeddedDocument)): if not any(isinstance(value, c) for c in choice_list): self.error( @@ -204,7 +208,8 @@ class BaseField(object): six.text_type(choice_list) ) ) - # Choices which are types other than Documents + + # Validate any other type of choices elif value not in choice_list: self.error('Value must be one of %s' % six.text_type(choice_list)) diff --git a/tests/fields/fields.py b/tests/fields/fields.py index 5c83d58c..50a88838 100644 --- a/tests/fields/fields.py +++ b/tests/fields/fields.py @@ -3205,8 +3205,22 @@ class FieldTest(unittest.TestCase): shirt.size = "XS" self.assertRaises(ValidationError, shirt.validate) + def test_choices_as_set(self): + """Ensure that sets can be used as field choices""" + class Shirt(Document): + size = StringField(choices={'S', 'M', 'L', 'XL', 'XXL'}) + Shirt.drop_collection() + shirt = Shirt() + shirt.validate() + + shirt.size = "S" + shirt.validate() + + shirt.size = "XS" + self.assertRaises(ValidationError, shirt.validate) + def test_choices_validation_documents(self): """ Ensure fields with document choices validate given a valid choice.