Merge pull request #2722 from bagerard/clone_list_enum_fields

[Clone] Fixing loading of EnumField inside ListField

Author: bagerard

docs/changelog.rst

@@ -18,6 +18,7 @@ Development
 - Addition of Decimal128Field: :class:`~mongoengine.fields.Decimal128Field` for accurate representation of Decimals (much better than the legacy field DecimalField).
   Although it could work to switch an existing DecimalField to Decimal128Field without applying a migration script,
   it is not recommended to do so (DecimalField uses float/str to store the value, Decimal128Field uses Decimal128).
+- BREAKING CHANGE: When using ListField(EnumField) or DictField(EnumField), the values weren't always cast into the Enum (#2531)
 
 Changes in 0.25.0
 =================

mongoengine/base/fields.py

@@ -282,6 +282,18 @@ def _lazy_load_refs(instance, name, ref_values, *, max_depth):
         )
         return documents
 
+    def __set__(self, instance, value):
+        # Some fields e.g EnumField are converted upon __set__
+        # So it is fair to mimic the same behavior when using e.g ListField(EnumField)
+        EnumField = _import_class("EnumField")
+        if self.field and isinstance(self.field, EnumField):
+            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:
@@ -434,12 +446,12 @@ def to_mongo(self, value, use_db_field=True, fields=None):
                             " have been saved to the database"
                         )
 
-                    # If its a document that is not inheritable it won't have
+                    # If it's a document that is not inheritable it won't have
                     # any _cls data so make it a generic reference allows
                     # 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

@@ -1122,7 +1122,7 @@ class MapField(DictField):
     """
 
     def __init__(self, field=None, *args, **kwargs):
-        # XXX ValidationError raised outside of the "validate" method.
+        # XXX ValidationError raised outside the "validate" method.
         if not isinstance(field, BaseField):
             self.error("Argument to MapField constructor must be a valid field")
         super().__init__(field=field, *args, **kwargs)
@@ -1665,14 +1665,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_binary_field.py

@@ -31,6 +31,14 @@ class Attachment(Document):
         assert MIME_TYPE == attachment_1.content_type
         assert BLOB == bytes(attachment_1.blob)
 
+    def test_bytearray_conversion_to_bytes(self):
+        class Dummy(Document):
+            blob = BinaryField()
+
+        byte_arr = bytearray(b"\x00\x00\x00\x00\x00")
+        dummy = Dummy(blob=byte_arr)
+        assert isinstance(dummy.blob, bytes)
+
     def test_validation_succeeds(self):
         """Ensure that valid values can be assigned to binary fields."""
 

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,42 @@ 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"]
+        model.save()
+        assert model.status == Status.DONE
+        assert model.statuses == [Status.NEW, Status.DONE]
+        assert model.color_mapping == {"blue": Color.BLUE}
+
+        model.reload()
+        assert model.status == Status.DONE
+        assert model.color_mapping == {"blue": Color.BLUE}
+        assert model.statuses == [Status.NEW, Status.DONE]
+
+        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 +172,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()
 
 

tests/fields/test_fields.py

@@ -44,6 +44,46 @@
 
 
 class TestField(MongoDBTestCase):
+    def test_constructor_set_historical_behavior_is_kept(self):
+        class MyDoc(Document):
+            oid = ObjectIdField()
+
+        doc = MyDoc()
+        doc.oid = str(ObjectId())
+        assert isinstance(doc.oid, str)
+
+        # not modified on save (historical behavior)
+        doc.save()
+        assert isinstance(doc.oid, str)
+
+        # reloading goes through constructor so it is expected to go through to_python
+        doc.reload()
+        assert isinstance(doc.oid, ObjectId)
+
+    def test_constructor_set_list_field_historical_behavior_is_kept(self):
+        # Although the behavior is not consistent between regular field and a ListField
+        # This is the historical behavior so we must make sure we don't modify it (unless consciously done of course)
+
+        class MyOIDSDoc(Document):
+            oids = ListField(ObjectIdField())
+
+        # constructor goes through to_python so casting occurs
+        doc = MyOIDSDoc(oids=[str(ObjectId())])
+        assert isinstance(doc.oids[0], ObjectId)
+
+        # constructor goes through to_python so casting occurs
+        doc = MyOIDSDoc()
+        doc.oids = [str(ObjectId())]
+        assert isinstance(doc.oids[0], str)
+
+        doc.save()
+        assert isinstance(doc.oids[0], str)
+
+        # reloading goes through constructor so it is expected to go through to_python
+        # and cast
+        doc.reload()
+        assert isinstance(doc.oids[0], ObjectId)
+
     def test_default_values_nothing_set(self):
         """Ensure that default field values are used when creating
         a document.

tests/fields/test_object_id_field.py

@@ -0,0 +1,30 @@
+import pytest
+from bson import ObjectId
+
+from mongoengine import Document, ObjectIdField, ValidationError
+from tests.utils import MongoDBTestCase, get_as_pymongo
+
+
+class TestObjectIdField(MongoDBTestCase):
+    def test_storage(self):
+        class MyDoc(Document):
+            oid = ObjectIdField()
+
+        doc = MyDoc(oid=ObjectId())
+        doc.save()
+        assert get_as_pymongo(doc) == {"_id": doc.id, "oid": doc.oid}
+
+    def test_constructor_converts_str_to_ObjectId(self):
+        class MyDoc(Document):
+            oid = ObjectIdField()
+
+        doc = MyDoc(oid=str(ObjectId()))
+        assert isinstance(doc.oid, ObjectId)
+
+    def test_validation_works(self):
+        class MyDoc(Document):
+            oid = ObjectIdField()
+
+        doc = MyDoc(oid="not-an-oid!")
+        with pytest.raises(ValidationError, match="Invalid ObjectID"):
+            doc.save()