From 9e40f3ae8328e2f755a8c681c5ff892d4e462d90 Mon Sep 17 00:00:00 2001 From: Mateusz Stankiewicz Date: Sat, 31 Oct 2020 10:47:20 +0100 Subject: [PATCH] PR ammends --- mongoengine/fields.py | 38 ++++++--- tests/fields/test_enum_field.py | 146 ++++++++++++++++---------------- 2 files changed, 99 insertions(+), 85 deletions(-) diff --git a/mongoengine/fields.py b/mongoengine/fields.py index 1ce66055..69277d06 100644 --- a/mongoengine/fields.py +++ b/mongoengine/fields.py @@ -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: diff --git a/tests/fields/test_enum_field.py b/tests/fields/test_enum_field.py index 00404aa5..1f89b9bf 100644 --- a/tests/fields/test_enum_field.py +++ b/tests/fields/test_enum_field.py @@ -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()