diff --git a/docs/changelog.rst b/docs/changelog.rst index 0e56baef..5e3107cf 100644 --- a/docs/changelog.rst +++ b/docs/changelog.rst @@ -9,6 +9,7 @@ Development - (Fill this out as you fix issues and develop your features). - 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 ================= diff --git a/mongoengine/fields.py b/mongoengine/fields.py index a3e483a4..7b2fe47f 100644 --- a/mongoengine/fields.py +++ b/mongoengine/fields.py @@ -1636,12 +1636,15 @@ class EnumField(BaseField): 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): diff --git a/tests/fields/test_enum_field.py b/tests/fields/test_enum_field.py index 7dfcffef..6dffb905 100644 --- a/tests/fields/test_enum_field.py +++ b/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 @@ class TestStringEnumField(MongoDBTestCase): 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"): + def test_partial_choices(self): + partial = [Status.DONE] + assert EnumField(Status, choices=partial).choices == partial + + def test_wrong_choices(self): + with pytest.raises(ValueError, match="Invalid choices"): EnumField(Status, choices=["my", "custom", "options"]) - - -class Color(Enum): - RED = 1 - BLUE = 2 + 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):