Merge pull request #2404 from mas15/add-enum-field

Add EnumField

AUTHORS

@@ -257,3 +257,4 @@ that much better:
  * Matthew Simpson (https://github.com/mcsimps2)
  * Leonardo Domingues (https://github.com/leodmgs)
  * Agustin Barto (https://github.com/abarto)
+ * Stankiewicz Mateusz (https://github.com/mas15)

docs/changelog.rst

@@ -13,6 +13,7 @@ Development
 - Fix the behavior of Doc.objects.limit(0) which should return all documents (similar to mongodb) #2311
 - Bug fix in ListField when updating the first item, it was saving the whole list, instead of
     just replacing the first item (as it's usually done) #2392
+- Add EnumField: ``mongoengine.fields.EnumField``
 
 Changes in 0.20.0
 =================

docs/guide/defining-documents.rst

@@ -76,6 +76,7 @@ are as follows:
 * :class:`~mongoengine.fields.EmailField`
 * :class:`~mongoengine.fields.EmbeddedDocumentField`
 * :class:`~mongoengine.fields.EmbeddedDocumentListField`
+* :class:`~mongoengine.fields.EnumField`
 * :class:`~mongoengine.fields.FileField`
 * :class:`~mongoengine.fields.FloatField`
 * :class:`~mongoengine.fields.GenericEmbeddedDocumentField`

mongoengine/fields.py

@@ -87,6 +87,7 @@
     "PolygonField",
     "SequenceField",
     "UUIDField",
+    "EnumField",
     "MultiPointField",
     "MultiLineStringField",
     "MultiPolygonField",
@@ -847,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
@@ -1622,6 +1622,68 @@ def prepare_query_value(self, op, value):
         return super().prepare_query_value(op, self.to_mongo(value))
 
 
+class EnumField(BaseField):
+    """Enumeration Field. Values are stored underneath as strings.
+     Example usage:
+    .. code-block:: python
+
+        class Status(Enum):
+            NEW = 'new'
+            DONE = 'done'
+
+        class ModelWithEnum(Document):
+            status = EnumField(Status, default=Status.NEW)
+
+        ModelWithEnum(status='done')
+        ModelWithEnum(status=Status.DONE)
+
+    Enum fields can be searched using enum or its value:
+    .. code-block:: python
+
+        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):
+        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 value.value
+        return value
+
+    def validate(self, value):
+        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:
+            return value
+        return super().prepare_query_value(op, self.to_mongo(value))
+
+
 class GridFSError(Exception):
     pass
 

tests/fields/test_enum_field.py

@@ -0,0 +1,107 @@
+from enum import Enum
+
+import pytest
+
+from mongoengine import *
+from tests.utils import MongoDBTestCase, get_as_pymongo
+
+
+class Status(Enum):
+    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
+
+
+class ModelWithColor(Document):
+    color = EnumField(Color, default=Color.RED)
+
+
+class TestIntEnumField(MongoDBTestCase):
+    def test_enum_with_int(self):
+        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}
+
+    def test_validate_model(self):
+        with pytest.raises(ValidationError, match="Value must be one of"):
+            ModelWithColor(color=3).validate()
+
+        with pytest.raises(ValidationError, match="Value must be one of"):
+            ModelWithColor(color="wrong_type").validate()