Allow to narrow down choices in an Enum field (#31)

Author: wojcikstefan

cleancat/base.py

@@ -1,4 +1,5 @@
 import datetime
+import inspect
 import re
 import sys
 
@@ -574,12 +575,31 @@ def clean(self, value):
 class Enum(Choices):
     """Like Choices, but expects a Python 3 Enum."""
 
+    def __init__(self, choices, **kwargs):
+        """Initialize the Enum field.
+
+        The `choices` param can be either:
+        * an enum.Enum class (in which case all of its values will become
+          valid choices),
+        * a list containing a subset of the enum's choices (e.g.
+          `[SomeEnumCls.OptionA, SomeEnumCls.OptionB]`). You must provide
+          more than one choice in this list and *all* of the choices must
+          belong to the same enum class.
+        """
+        is_cls = inspect.isclass(choices)
+        if is_cls:
+            self.enum_cls = choices
+        else:
+            assert choices, 'You need to provide at least one enum choice.'
+            self.enum_cls = choices[0].__class__
+        return super(Enum, self).__init__(choices, **kwargs)
+
     def get_choices(self):
         return [choice.value for choice in self.choices]
 
     def clean(self, value):
         value = super(Enum, self).clean(value)
-        return self.choices(value)
+        return self.enum_cls(value)
 
     def serialize(self, choice):
         if choice is not None:

setup.py

@@ -1,3 +1,4 @@
+import sys
 from setuptools import setup
 
 install_requirements = [
@@ -11,6 +12,9 @@
     'sqlalchemy'
 ]
 
+if sys.version_info[:2] < (3, 4):
+    test_requirements += ['enum34']
+
 setup(
     name='cleancat',
     version='0.7.3',

tests/__init__.py

@@ -1,7 +1,7 @@
 # -*- coding: utf-8 -*-
 import datetime
+import enum
 import re
-import sys
 import unittest
 
 import pytest
@@ -319,10 +319,7 @@ class CaseInsensitiveChoiceSchema(Schema):
         self.assertInvalid(CaseInsensitiveChoiceSchema({'choice': 'world '}), {'field-errors': ['choice']})
         self.assertInvalid(CaseInsensitiveChoiceSchema({'choice': 'invalid'}), {'field-errors': ['choice']})
 
-    @unittest.skipIf(sys.version_info < (3, 4), 'enum unavailable')
     def test_enum(self):
-        import enum
-
         class MyChoices(enum.Enum):
             A = 'a'
             B = 'b'
@@ -334,6 +331,19 @@ class ChoiceSchema(Schema):
         self.assertValid(ChoiceSchema({'choice': 'b'}), {'choice': MyChoices.B})
         self.assertInvalid(ChoiceSchema({'choice': 'c'}), {'field-errors': ['choice']})
 
+    def test_enum_with_choices(self):
+        class MyChoices(enum.Enum):
+            A = 'a'
+            B = 'b'
+            C = 'c'
+
+        class ChoiceSchema(Schema):
+            choice = Enum([MyChoices.A, MyChoices.B])
+
+        self.assertValid(ChoiceSchema({'choice': 'a'}), {'choice': MyChoices.A})
+        self.assertValid(ChoiceSchema({'choice': 'b'}), {'choice': MyChoices.B})
+        self.assertInvalid(ChoiceSchema({'choice': 'c'}), {'field-errors': ['choice']})
+
     def test_url(self):
         class URLSchema(Schema):
             url = URL()
@@ -658,10 +668,7 @@ class TestSchema(Schema):
             'dictionary': {},
         }
 
-    @unittest.skipIf(sys.version_info < (3, 4), 'enum unavailable')
     def test_serialization_enum(self):
-        import enum
-
         class MyChoices(enum.Enum):
             A = 'a'
             B = 'b'