PR ammends
This commit is contained in:
Mateusz Stankiewicz
parent
c9d53ca5d5
commit
9e40f3ae83
@ -848,8 +848,7 @@ class DynamicField(BaseField):
|
||||
Used by :class:`~mongoengine.DynamicDocument` to handle dynamic data"""
|
||||
|
||||
def to_mongo(self, value, use_db_field=True, fields=None):
|
||||
"""Convert a Python type to a MongoDB compatible type.
|
||||
"""
|
||||
"""Convert a Python type to a MongoDB compatible type."""
|
||||
|
||||
if isinstance(value, str):
|
||||
return value
|
||||
@ -1624,7 +1623,7 @@ class BinaryField(BaseField):
|
||||
|
||||
|
||||
class EnumField(BaseField):
|
||||
""" Enumeration Field. Values are stored underneath as strings.
|
||||
"""Enumeration Field. Values are stored underneath as strings.
|
||||
Example usage:
|
||||
.. code-block:: python
|
||||
|
||||
@ -1643,30 +1642,41 @@ class EnumField(BaseField):
|
||||
|
||||
ModelWithEnum.objects(status='new').count()
|
||||
ModelWithEnum.objects(status=Status.NEW).count()
|
||||
|
||||
Note that choices cannot be set explicitly, they are derived
|
||||
from the provided enum class.
|
||||
"""
|
||||
|
||||
def __init__(self, enum, **kwargs):
|
||||
self._enum_cls = enum
|
||||
if "choices" in kwargs:
|
||||
raise ValueError(
|
||||
"'choices' can't be set on EnumField, "
|
||||
"it is implicitly set as the enum class"
|
||||
)
|
||||
kwargs["choices"] = list(self._enum_cls)
|
||||
super().__init__(**kwargs)
|
||||
|
||||
def __set__(self, instance, value):
|
||||
if value is None or isinstance(value, self._enum_cls):
|
||||
value = value # if it is proper enum or none then fine
|
||||
else: # if it not, then try to create enum of it
|
||||
value = self._enum_cls(value)
|
||||
is_legal_value = value is None or isinstance(value, self._enum_cls)
|
||||
if not is_legal_value:
|
||||
try:
|
||||
value = self._enum_cls(value)
|
||||
except Exception:
|
||||
pass
|
||||
return super().__set__(instance, value)
|
||||
|
||||
def to_mongo(self, value):
|
||||
if isinstance(value, self._enum_cls):
|
||||
return str(value.value)
|
||||
return str(value)
|
||||
return value.value
|
||||
return value
|
||||
|
||||
def validate(self, value):
|
||||
if not isinstance(value, self._enum_cls):
|
||||
self.error(
|
||||
"EnumField only accepts instances of "
|
||||
"(%s)" % self._enum_cls
|
||||
)
|
||||
if value and not isinstance(value, self._enum_cls):
|
||||
try:
|
||||
self._enum_cls(value)
|
||||
except Exception as e:
|
||||
self.error(str(e))
|
||||
|
||||
def prepare_query_value(self, op, value):
|
||||
if value is None:
|
||||
|
||||
@ -1,4 +1,3 @@
|
||||
# -*- coding: utf-8 -*-
|
||||
from enum import Enum
|
||||
|
||||
import pytest
|
||||
@ -8,14 +7,72 @@ from tests.utils import MongoDBTestCase, get_as_pymongo
|
||||
|
||||
|
||||
class Status(Enum):
|
||||
NEW = 'new'
|
||||
DONE = 'done'
|
||||
|
||||
|
||||
NEW = "new"
|
||||
DONE = "done"
|
||||
|
||||
|
||||
class ModelWithEnum(Document):
|
||||
status = EnumField(Status)
|
||||
|
||||
|
||||
class TestStringEnumField(MongoDBTestCase):
|
||||
def test_storage(self):
|
||||
model = ModelWithEnum(status=Status.NEW).save()
|
||||
assert get_as_pymongo(model) == {"_id": model.id, "status": "new"}
|
||||
|
||||
def test_set_enum(self):
|
||||
ModelWithEnum.drop_collection()
|
||||
ModelWithEnum(status=Status.NEW).save()
|
||||
assert ModelWithEnum.objects(status=Status.NEW).count() == 1
|
||||
assert ModelWithEnum.objects.first().status == Status.NEW
|
||||
|
||||
def test_set_by_value(self):
|
||||
ModelWithEnum.drop_collection()
|
||||
ModelWithEnum(status="new").save()
|
||||
assert ModelWithEnum.objects.first().status == Status.NEW
|
||||
|
||||
def test_filter(self):
|
||||
ModelWithEnum.drop_collection()
|
||||
ModelWithEnum(status="new").save()
|
||||
assert ModelWithEnum.objects(status="new").count() == 1
|
||||
assert ModelWithEnum.objects(status=Status.NEW).count() == 1
|
||||
assert ModelWithEnum.objects(status=Status.DONE).count() == 0
|
||||
|
||||
def test_change_value(self):
|
||||
m = ModelWithEnum(status="new")
|
||||
m.status = Status.DONE
|
||||
m.save()
|
||||
assert m.status == Status.DONE
|
||||
|
||||
def test_set_default(self):
|
||||
class ModelWithDefault(Document):
|
||||
status = EnumField(Status, default=Status.DONE)
|
||||
|
||||
m = ModelWithDefault().save()
|
||||
assert m.status == Status.DONE
|
||||
|
||||
def test_enum_field_can_be_empty(self):
|
||||
ModelWithEnum.drop_collection()
|
||||
m = ModelWithEnum().save()
|
||||
assert m.status is None
|
||||
assert ModelWithEnum.objects()[0].status is None
|
||||
assert ModelWithEnum.objects(status=None).count() == 1
|
||||
|
||||
def test_set_none_explicitly(self):
|
||||
ModelWithEnum.drop_collection()
|
||||
ModelWithEnum(status=None).save()
|
||||
assert ModelWithEnum.objects.first().status is None
|
||||
|
||||
def test_cannot_create_model_with_wrong_enum_value(self):
|
||||
m = ModelWithEnum(status="wrong_one")
|
||||
with pytest.raises(ValidationError):
|
||||
m.validate()
|
||||
|
||||
def test_user_is_informed_when_tries_to_set_choices(self):
|
||||
with pytest.raises(ValueError, match="'choices' can't be set on EnumField"):
|
||||
EnumField(Status, choices=["my", "custom", "options"])
|
||||
|
||||
|
||||
class Color(Enum):
|
||||
RED = 1
|
||||
BLUE = 2
|
||||
@ -25,79 +82,26 @@ class ModelWithColor(Document):
|
||||
color = EnumField(Color, default=Color.RED)
|
||||
|
||||
|
||||
class TestEnumField(MongoDBTestCase):
|
||||
def test_storage(self):
|
||||
model = ModelWithEnum(status=Status.NEW).save()
|
||||
assert get_as_pymongo(model) == {"_id": model.id, "status": 'new'}
|
||||
|
||||
def test_set_enum(self):
|
||||
ModelWithEnum.drop_collection()
|
||||
m = ModelWithEnum(status=Status.NEW).save()
|
||||
assert ModelWithEnum.objects(status=Status.NEW).count() == 1
|
||||
assert ModelWithEnum.objects.first().status == Status.NEW
|
||||
m.validate()
|
||||
|
||||
def test_set_by_value(self):
|
||||
ModelWithEnum.drop_collection()
|
||||
ModelWithEnum(status='new').save()
|
||||
assert ModelWithEnum.objects.first().status == Status.NEW
|
||||
|
||||
def test_filter(self):
|
||||
ModelWithEnum.drop_collection()
|
||||
ModelWithEnum(status='new').save()
|
||||
assert ModelWithEnum.objects(status='new').count() == 1
|
||||
assert ModelWithEnum.objects(status=Status.NEW).count() == 1
|
||||
assert ModelWithEnum.objects(status=Status.DONE).count() == 0
|
||||
|
||||
def test_change_value(self):
|
||||
m = ModelWithEnum(status='new')
|
||||
m.status = Status.DONE
|
||||
m.validate()
|
||||
assert m.status == Status.DONE
|
||||
|
||||
def test_set_default(self):
|
||||
class ModelWithDefault(Document):
|
||||
status = EnumField(Status, default=Status.DONE)
|
||||
|
||||
m = ModelWithDefault()
|
||||
m.validate()
|
||||
m.save()
|
||||
assert m.status == Status.DONE
|
||||
|
||||
class TestIntEnumField(MongoDBTestCase):
|
||||
def test_enum_with_int(self):
|
||||
m = ModelWithColor()
|
||||
m.validate()
|
||||
m.save()
|
||||
ModelWithColor.drop_collection()
|
||||
m = ModelWithColor().save()
|
||||
assert m.color == Color.RED
|
||||
assert ModelWithColor.objects(color=Color.RED).count() == 1
|
||||
assert ModelWithColor.objects(color=1).count() == 1
|
||||
assert ModelWithColor.objects(color=2).count() == 0
|
||||
|
||||
def test_create_int_enum_by_value(self):
|
||||
model = ModelWithColor(color=2).save()
|
||||
assert model.color == Color.BLUE
|
||||
|
||||
def test_storage_enum_with_int(self):
|
||||
model = ModelWithColor(color=Color.BLUE).save()
|
||||
assert get_as_pymongo(model) == {"_id": model.id, "color": "2"}
|
||||
assert get_as_pymongo(model) == {"_id": model.id, "color": 2}
|
||||
|
||||
def test_enum_field_can_be_empty(self):
|
||||
m = ModelWithEnum()
|
||||
m.validate()
|
||||
m.save()
|
||||
assert m.status is None
|
||||
assert ModelWithEnum.objects()[0].status is None
|
||||
assert ModelWithEnum.objects(status=None).count() == 1
|
||||
def test_validate_model(self):
|
||||
with pytest.raises(ValidationError, match="Value must be one of"):
|
||||
ModelWithColor(color=3).validate()
|
||||
|
||||
def test_cannot_create_model_with_wrong_enum_value(self):
|
||||
with pytest.raises(ValueError):
|
||||
ModelWithEnum(status='wrong_one')
|
||||
|
||||
def test_cannot_create_model_with_wrong_enum_type(self):
|
||||
with pytest.raises(ValueError):
|
||||
ModelWithColor(color='wrong_type')
|
||||
|
||||
def test_cannot_create_model_with_wrong_enum_value_int(self):
|
||||
with pytest.raises(ValueError):
|
||||
ModelWithColor(color=3)
|
||||
|
||||
def test_cannot_set_wrong_enum_value(self):
|
||||
m = ModelWithEnum(status='new')
|
||||
with pytest.raises(ValueError):
|
||||
m.status = 'wrong'
|
||||
with pytest.raises(ValidationError, match="Value must be one of"):
|
||||
ModelWithColor(color="wrong_type").validate()
|
||||
|
||||
Loading…
x
Reference in New Issue
Block a user