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 @@ __all__ = (
     "PolygonField",
     "SequenceField",
     "UUIDField",
+    "EnumField",
     "MultiPointField",
     "MultiLineStringField",
     "MultiPolygonField",
@@ -1622,6 +1623,57 @@ class BinaryField(BaseField):
         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()
+    """
+    def __init__(self, enum, **kwargs):
+        self._enum_cls = enum
+        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)
+        return super().__set__(instance, value)
+
+    def to_mongo(self, value):
+        if isinstance(value, self._enum_cls):
+            return str(value.value)
+        return str(value)
+
+    def validate(self, value):
+        if not isinstance(value, self._enum_cls):
+            self.error(
+                "EnumField only accepts instances of "
+                "(%s)" % self._enum_cls
+            )
+
+    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,103 @@
+# -*- coding: utf-8 -*-
+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 Color(Enum):
+    RED = 1
+    BLUE = 2
+
+
+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
+
+    def test_enum_with_int(self):
+        m = ModelWithColor()
+        m.validate()
+        m.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_storage_enum_with_int(self):
+        model = ModelWithColor(color=Color.BLUE).save()
+        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_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'