Merge pull request #2509 from jaesivsm/fix/EnumField-choices

Fixing EnumField choices validation

docs/changelog.rst

@@ -12,6 +12,7 @@ Changes in 0.23.1
 ===========
 - Bug fix: ignore LazyReferenceFields when clearing _changed_fields #2484
 - Improve connection doc #2481
+- EnumField improvements: now `choices` limits the values of an enum to allow
 
 Changes in 0.23.0
 =================

mongoengine/fields.py

@@ -1636,12 +1636,15 @@ class ModelWithEnum(Document):
 
     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)  # Implicit validator
+        if kwargs.get("choices"):
+            invalid_choices = []
+            for choice in kwargs["choices"]:
+                if not isinstance(choice, enum):
+                    invalid_choices.append(choice)
+            if invalid_choices:
+                raise ValueError("Invalid choices: %r" % invalid_choices)
+        else:
+            kwargs["choices"] = list(self._enum_cls)  # Implicit validator
         super().__init__(**kwargs)
 
     def __set__(self, instance, value):

tests/fields/test_enum_field.py

@@ -12,6 +12,11 @@ class Status(Enum):
     DONE = "done"
 
 
+class Color(Enum):
+    RED = 1
+    BLUE = 2
+
+
 class ModelWithEnum(Document):
     status = EnumField(Status)
 
@@ -74,14 +79,17 @@ def test_cannot_create_model_with_wrong_enum_value(self):
         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"])
+    def test_partial_choices(self):
+        partial = [Status.DONE]
+        assert EnumField(Status, choices=partial).choices == partial
 
-
-class Color(Enum):
-    RED = 1
-    BLUE = 2
+    def test_wrong_choices(self):
+        with pytest.raises(ValueError, match="Invalid choices"):
+            EnumField(Status, choices=["my", "custom", "options"])
+        with pytest.raises(ValueError, match="Invalid choices"):
+            EnumField(Status, choices=[Color.RED])
+        with pytest.raises(ValueError, match="Invalid choices"):
+            EnumField(Status, choices=[Status.DONE, Color.RED])
 
 
 class ModelWithColor(Document):