pub/mongoengine
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Add choices keyword argument to BaseField.__init__()

Browse Source
This commit is contained in:
Don Spaulding 2010-04-15 17:59:35 -05:00
parent e9c92f30ba
commit ee0c75a26d
2 changed files with 32 additions and 3 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

14
mongoengine/base.py
View File

@ -24,7 +24,8 @@ class BaseField(object):
_index_with_types = True _index_with_types = True
def __init__(self, db_field=None, name=None, required=False, default=None, def __init__(self, db_field=None, name=None, required=False, default=None,
unique=False, unique_with=None, primary_key=False): unique=False, unique_with=None, primary_key=False,
choices=None):
self.db_field = (db_field or name) if not primary_key else '_id' self.db_field = (db_field or name) if not primary_key else '_id'
if name: if name:
import warnings import warnings
@ -36,6 +37,7 @@ class BaseField(object):
self.unique = bool(unique or unique_with) self.unique = bool(unique or unique_with)
self.unique_with = unique_with self.unique_with = unique_with
self.primary_key = primary_key self.primary_key = primary_key
self.choices = choices
def __get__(self, instance, owner): def __get__(self, instance, owner):
"""Descriptor for retrieving a value from a field in a document. Do """Descriptor for retrieving a value from a field in a document. Do
@ -79,6 +81,12 @@ class BaseField(object):
""" """
pass pass
def _validate(self, value):
if self.choices is not None:
if value not in self.choices:
raise ValidationError("Value must be one of %s."%unicode(self.choices))
self.validate(value)
class ObjectIdField(BaseField): class ObjectIdField(BaseField):
"""An field wrapper around MongoDB's ObjectIds. """An field wrapper around MongoDB's ObjectIds.
@ -314,7 +322,7 @@ class BaseDocument(object):
for field, value in fields: for field, value in fields:
if value is not None: if value is not None:
try: try:
field.validate(value) field._validate(value)
except (ValueError, AttributeError, AssertionError), e: except (ValueError, AttributeError, AssertionError), e:
raise ValidationError('Invalid value for field of type "' + raise ValidationError('Invalid value for field of type "' +
field.__class__.__name__ + '"') field.__class__.__name__ + '"')
@ -441,4 +449,4 @@ if sys.version_info < (2, 5):
return types.ClassType(name, parents, {}) return types.ClassType(name, parents, {})
else: else:
def subclass_exception(name, parents, module): def subclass_exception(name, parents, module):
return type(name, parents, {'__module__': module}) return type(name, parents, {'__module__': module})

21
tests/fields.py
View File

@ -588,5 +588,26 @@ class FieldTest(unittest.TestCase):
AttachmentRequired.drop_collection() AttachmentRequired.drop_collection()
AttachmentSizeLimit.drop_collection() AttachmentSizeLimit.drop_collection()
def test_choices_validation(self):
"""Ensure that value is in a container of allowed values.
"""
class Shirt(Document):
size = StringField(max_length=3, 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)
Shirt.drop_collection()
if __name__ == '__main__': if __name__ == '__main__':
unittest.main() unittest.main()