Fixing loading of EnumField inside ListField

Under certain condition, inside value of ListField(EnumField) don't get casted to the proper Enum. This patches fixes that.

Author: jaesivsm

mongoengine/base/fields.py

@@ -134,9 +134,7 @@ def __set__(self, instance, value):
         # If setting to None and there is a default value provided for this
         # field, then set the value to the default value.
         if value is None:
-            if self.null:
-                value = None
-            elif self.default is not None:
+            if not self.null and self.default is not None:
                 value = self.default
                 if callable(value):
                     value = value()
@@ -282,6 +280,19 @@ def _lazy_load_refs(instance, name, ref_values, *, max_depth):
         )
         return documents
 
+    def __set__(self, instance, value):
+        unprocessable_fields = (
+            ComplexBaseField,
+            _import_class("EmbeddedDocumentField"),
+            _import_class("FileField"),
+        )
+        if self.field and not isinstance(self.field, unprocessable_fields):
+            if isinstance(value, (list, tuple)):
+                value = [self.field.to_python(sub_val) for sub_val in value]
+            elif isinstance(value, dict):
+                value = {key: self.field.to_python(sub) for key, sub in value.items()}
+        return super().__set__(instance, value)
+
     def __get__(self, instance, owner):
         """Descriptor to automatically dereference references."""
         if instance is None:
@@ -439,7 +450,7 @@ def to_mongo(self, value, use_db_field=True, fields=None):
                     # us to dereference
                     meta = getattr(v, "_meta", {})
                     allow_inheritance = meta.get("allow_inheritance")
-                    if not allow_inheritance and not self.field:
+                    if not allow_inheritance:
                         value_dict[k] = GenericReferenceField().to_mongo(v)
                     else:
                         collection = v._get_collection_name()

mongoengine/fields.py

@@ -1660,14 +1660,25 @@ def __init__(self, enum, **kwargs):
             kwargs["choices"] = list(self._enum_cls)  # Implicit validator
         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:
+    def validate(self, value):
+        if isinstance(value, self._enum_cls):
+            return super().validate(value)
+        try:
+            self._enum_cls(value)
+        except ValueError:
+            self.error(f"{value} is not a valid {self._enum_cls}")
+
+    def to_python(self, value):
+        value = super().to_python(value)
+        if not isinstance(value, self._enum_cls):
             try:
-                value = self._enum_cls(value)
-            except Exception:
-                pass
-        return super().__set__(instance, value)
+                return self._enum_cls(value)
+            except ValueError:
+                return value
+        return value
+
+    def __set__(self, instance, value):
+        return super().__set__(instance, self.to_python(value))
 
     def to_mongo(self, value):
         if isinstance(value, self._enum_cls):

tests/fields/test_enum_field.py

@@ -3,7 +3,13 @@
 import pytest
 from bson import InvalidDocument
 
-from mongoengine import Document, EnumField, ValidationError
+from mongoengine import (
+    DictField,
+    Document,
+    EnumField,
+    ListField,
+    ValidationError,
+)
 from tests.utils import MongoDBTestCase, get_as_pymongo
 
 
@@ -21,6 +27,12 @@ class ModelWithEnum(Document):
     status = EnumField(Status)
 
 
+class ModelComplexEnum(Document):
+    status = EnumField(Status)
+    statuses = ListField(EnumField(Status))
+    color_mapping = DictField(EnumField(Color))
+
+
 class TestStringEnumField(MongoDBTestCase):
     def test_storage(self):
         model = ModelWithEnum(status=Status.NEW).save()
@@ -101,6 +113,38 @@ def test_wrong_choices(self):
         with pytest.raises(ValueError, match="Invalid choices"):
             EnumField(Status, choices=[Status.DONE, Color.RED])
 
+    def test_embedding_in_complex_field(self):
+        ModelComplexEnum.drop_collection()
+        model = ModelComplexEnum(
+            status="new", statuses=["new"], color_mapping={"red": 1}
+        ).save()
+        assert model.status == Status.NEW
+        assert model.statuses == [Status.NEW]
+        assert model.color_mapping == {"red": Color.RED}
+        model.reload()
+        assert model.status == Status.NEW
+        assert model.statuses == [Status.NEW]
+        assert model.color_mapping == {"red": Color.RED}
+        model.status = "done"
+        model.color_mapping = {"blue": 2}
+        model.statuses = ["new", "done"]
+        assert model.status == Status.DONE
+        assert model.color_mapping == {"blue": Color.BLUE}, model.color_mapping
+        assert model.statuses == [Status.NEW, Status.DONE], model.statuses
+        model = model.save().reload()
+        assert model.status == Status.DONE
+        assert model.color_mapping == {"blue": Color.BLUE}, model.color_mapping
+        assert model.statuses == [Status.NEW, Status.DONE], model.statuses
+
+        with pytest.raises(ValidationError, match="must be one of ..Status"):
+            model.statuses = [1]
+            model.save()
+
+        model.statuses = ["done"]
+        model.color_mapping = {"blue": "done"}
+        with pytest.raises(ValidationError, match="must be one of ..Color"):
+            model.save()
+
 
 class ModelWithColor(Document):
     color = EnumField(Color, default=Color.RED)
@@ -124,10 +168,7 @@ def test_storage_enum_with_int(self):
         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"):
+        with pytest.raises(ValidationError, match="must be one of ..Color"):
             ModelWithColor(color="wrong_type").validate()