diff --git a/AUTHORS b/AUTHORS index 02e43955..10d04c68 100644 --- a/AUTHORS +++ b/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) diff --git a/docs/changelog.rst b/docs/changelog.rst index f616f4a6..f63edc61 100644 --- a/docs/changelog.rst +++ b/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 ================= diff --git a/docs/guide/defining-documents.rst b/docs/guide/defining-documents.rst index f5c70728..7fc20ba8 100644 --- a/docs/guide/defining-documents.rst +++ b/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` diff --git a/mongoengine/fields.py b/mongoengine/fields.py index c5926cbd..69277d06 100644 --- a/mongoengine/fields.py +++ b/mongoengine/fields.py @@ -87,6 +87,7 @@ __all__ = ( "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 @@ 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() + + 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 diff --git a/tests/fields/test_enum_field.py b/tests/fields/test_enum_field.py new file mode 100644 index 00000000..1f89b9bf --- /dev/null +++ b/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()