PR ammends
This commit is contained in:
parent
c9d53ca5d5
commit
9e40f3ae83
@ -848,8 +848,7 @@ class DynamicField(BaseField):
|
|||||||
Used by :class:`~mongoengine.DynamicDocument` to handle dynamic data"""
|
Used by :class:`~mongoengine.DynamicDocument` to handle dynamic data"""
|
||||||
|
|
||||||
def to_mongo(self, value, use_db_field=True, fields=None):
|
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):
|
if isinstance(value, str):
|
||||||
return value
|
return value
|
||||||
@ -1643,30 +1642,41 @@ class EnumField(BaseField):
|
|||||||
|
|
||||||
ModelWithEnum.objects(status='new').count()
|
ModelWithEnum.objects(status='new').count()
|
||||||
ModelWithEnum.objects(status=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):
|
def __init__(self, enum, **kwargs):
|
||||||
self._enum_cls = enum
|
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)
|
kwargs["choices"] = list(self._enum_cls)
|
||||||
super().__init__(**kwargs)
|
super().__init__(**kwargs)
|
||||||
|
|
||||||
def __set__(self, instance, value):
|
def __set__(self, instance, value):
|
||||||
if value is None or isinstance(value, self._enum_cls):
|
is_legal_value = value is None or isinstance(value, self._enum_cls)
|
||||||
value = value # if it is proper enum or none then fine
|
if not is_legal_value:
|
||||||
else: # if it not, then try to create enum of it
|
try:
|
||||||
value = self._enum_cls(value)
|
value = self._enum_cls(value)
|
||||||
|
except Exception:
|
||||||
|
pass
|
||||||
return super().__set__(instance, value)
|
return super().__set__(instance, value)
|
||||||
|
|
||||||
def to_mongo(self, value):
|
def to_mongo(self, value):
|
||||||
if isinstance(value, self._enum_cls):
|
if isinstance(value, self._enum_cls):
|
||||||
return str(value.value)
|
return value.value
|
||||||
return str(value)
|
return value
|
||||||
|
|
||||||
def validate(self, value):
|
def validate(self, value):
|
||||||
if not isinstance(value, self._enum_cls):
|
if value and not isinstance(value, self._enum_cls):
|
||||||
self.error(
|
try:
|
||||||
"EnumField only accepts instances of "
|
self._enum_cls(value)
|
||||||
"(%s)" % self._enum_cls
|
except Exception as e:
|
||||||
)
|
self.error(str(e))
|
||||||
|
|
||||||
def prepare_query_value(self, op, value):
|
def prepare_query_value(self, op, value):
|
||||||
if value is None:
|
if value is None:
|
||||||
|
@ -1,4 +1,3 @@
|
|||||||
# -*- coding: utf-8 -*-
|
|
||||||
from enum import Enum
|
from enum import Enum
|
||||||
|
|
||||||
import pytest
|
import pytest
|
||||||
@ -8,14 +7,72 @@ from tests.utils import MongoDBTestCase, get_as_pymongo
|
|||||||
|
|
||||||
|
|
||||||
class Status(Enum):
|
class Status(Enum):
|
||||||
NEW = 'new'
|
NEW = "new"
|
||||||
DONE = 'done'
|
DONE = "done"
|
||||||
|
|
||||||
|
|
||||||
class ModelWithEnum(Document):
|
class ModelWithEnum(Document):
|
||||||
status = EnumField(Status)
|
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):
|
class Color(Enum):
|
||||||
RED = 1
|
RED = 1
|
||||||
BLUE = 2
|
BLUE = 2
|
||||||
@ -25,79 +82,26 @@ class ModelWithColor(Document):
|
|||||||
color = EnumField(Color, default=Color.RED)
|
color = EnumField(Color, default=Color.RED)
|
||||||
|
|
||||||
|
|
||||||
class TestEnumField(MongoDBTestCase):
|
class TestIntEnumField(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
|
|
||||||
|
|
||||||
def test_enum_with_int(self):
|
def test_enum_with_int(self):
|
||||||
m = ModelWithColor()
|
ModelWithColor.drop_collection()
|
||||||
m.validate()
|
m = ModelWithColor().save()
|
||||||
m.save()
|
|
||||||
assert m.color == Color.RED
|
assert m.color == Color.RED
|
||||||
assert ModelWithColor.objects(color=Color.RED).count() == 1
|
assert ModelWithColor.objects(color=Color.RED).count() == 1
|
||||||
assert ModelWithColor.objects(color=1).count() == 1
|
assert ModelWithColor.objects(color=1).count() == 1
|
||||||
assert ModelWithColor.objects(color=2).count() == 0
|
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):
|
def test_storage_enum_with_int(self):
|
||||||
model = ModelWithColor(color=Color.BLUE).save()
|
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):
|
def test_validate_model(self):
|
||||||
m = ModelWithEnum()
|
with pytest.raises(ValidationError, match="Value must be one of"):
|
||||||
m.validate()
|
ModelWithColor(color=3).validate()
|
||||||
m.save()
|
|
||||||
assert m.status is None
|
|
||||||
assert ModelWithEnum.objects()[0].status is None
|
|
||||||
assert ModelWithEnum.objects(status=None).count() == 1
|
|
||||||
|
|
||||||
def test_cannot_create_model_with_wrong_enum_value(self):
|
with pytest.raises(ValidationError, match="Value must be one of"):
|
||||||
with pytest.raises(ValueError):
|
ModelWithColor(color="wrong_type").validate()
|
||||||
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'
|
|
||||||
|
Loading…
x
Reference in New Issue
Block a user