PR ammends

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:

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()