Merge pull request #91 from alfred82santa/feature/enum-field

Add new Enum field & tests

dirty_models/fields.py

@@ -11,7 +11,7 @@
 
 __all__ = ['IntegerField', 'FloatField', 'BooleanField', 'StringField', 'StringIdField',
            'TimeField', 'DateField', 'DateTimeField', 'TimedeltaField', 'ModelField', 'ArrayField',
-           'HashMapField', 'BlobField', 'MultiTypeField']
+           'HashMapField', 'BlobField', 'MultiTypeField', 'EnumField']
 
 
 class BaseField:
@@ -124,7 +124,7 @@ def check_value(self, value):
 
     def can_use_value(self, value):
         return isinstance(value, float) \
-               or (isinstance(value, str) and value.isdigit())
+            or (isinstance(value, str) and value.isdigit())
 
 
 class FloatField(BaseField):
@@ -146,7 +146,7 @@ def check_value(self, value):
 
     def can_use_value(self, value):
         return isinstance(value, int) or \
-               (isinstance(value, str) and
+            (isinstance(value, str) and
                 value.replace('.', '', 1).isnumeric())
 
 
@@ -607,6 +607,7 @@ def __set__(self, obj, value):
 
 
 class InnerFieldTypeMixin:
+
     def __init__(self, field_type=None, **kwargs):
         self._field_type = None
         if isinstance(field_type, tuple):
@@ -787,3 +788,55 @@ def get_field_type_by_value(self, value):
     @property
     def field_types(self):
         return self._field_types.copy()
+
+
+class EnumField(BaseField):
+    """
+    It allows to create a field which contains a member of an enumeration.
+
+    **Automatic cast from:**
+
+    * Any value of enumeration.
+
+    * Any member name of enumeration.
+    """
+
+    def __init__(self, enum_class, *args, **kwargs):
+        """
+
+        :param enum_class: Enumeration class
+        :type enum_class: enum.Enum
+        """
+        self.enum_class = enum_class
+        super(EnumField, self).__init__(*args, **kwargs)
+
+    def export_definition(self):
+        result = super(EnumField, self).export_definition()
+        result['enum_class'] = self.enum_class
+
+        return result
+
+    def get_field_docstring(self):
+        dcstr = super(EnumField, self).get_field_docstring()
+
+        if self.enum_class:
+            dcstr += ' (:class:`{0}`)'.format('.'.join([self.enum_class.__module__, self.enum_class.__name__]))
+        return dcstr
+
+    def convert_value(self, value):
+        try:
+            return self.enum_class(value)
+        except ValueError:
+            return getattr(self.enum_class, value)
+
+    def check_value(self, value):
+        return isinstance(value, self.enum_class)
+
+    def can_use_value(self, value):
+        try:
+            self.enum_class(value)
+            return True
+        except ValueError:
+            pass
+
+        return value in self.enum_class.__members__.keys()

tests/dirty_models/tests_fields.py

@@ -1,12 +1,14 @@
 from datetime import time, date, datetime, timezone, timedelta
+from enum import Enum
 from unittest import TestCase
 
 import iso8601
 from dateutil import tz
 
 from dirty_models.fields import (IntegerField, StringField, BooleanField,
                                  FloatField, ModelField, TimeField, DateField,
-                                 DateTimeField, ArrayField, StringIdField, HashMapField, MultiTypeField, TimedeltaField)
+                                 DateTimeField, ArrayField, StringIdField, HashMapField, MultiTypeField, TimedeltaField,
+                                 EnumField)
 from dirty_models.model_types import ListModel
 from dirty_models.models import BaseModel, HashMapModel
 
@@ -1531,7 +1533,7 @@ def test_export_definition(self):
                          field.export_definition())
 
 
-class TimeFieldWithTimezone(TestCase):
+class TimeFieldWithTimezoneTests(TestCase):
 
     def test_no_timezone_none(self):
         class Model(BaseModel):
@@ -1581,3 +1583,51 @@ def test_export_definition(self):
                           'default_timezone': timezone.utc,
                           'name': 'test_field', 'read_only': False},
                          field.export_definition())
+
+
+class EnumFieldTests(TestCase):
+
+    class TestEnum(Enum):
+
+        value_1 = 'value1'
+        value_2 = 2
+
+    def setUp(self):
+        self.field = EnumField(name='test_field', alias=[], enum_class=self.TestEnum)
+
+    def test_check_value(self):
+        self.assertTrue(self.field.check_value(self.TestEnum.value_1))
+        self.assertTrue(self.field.check_value(self.TestEnum.value_2))
+
+    def test_check_value_fail(self):
+        self.assertFalse(self.field.check_value('value_1'))
+        self.assertFalse(self.field.check_value(2))
+
+    def test_can_use_value_check_values(self):
+        self.assertTrue(self.field.can_use_value('value1'))
+        self.assertTrue(self.field.can_use_value(2))
+
+    def test_can_use_value_check_member_names(self):
+        self.assertTrue(self.field.can_use_value('value_1'))
+        self.assertTrue(self.field.can_use_value('value_2'))
+
+    def test_can_use_value_check_values_fail(self):
+        self.assertFalse(self.field.can_use_value('value2'))
+        self.assertFalse(self.field.can_use_value(3))
+
+    def test_convert_value_from_values(self):
+        self.assertEqual(self.field.convert_value('value1'), self.TestEnum.value_1)
+        self.assertEqual(self.field.convert_value(2), self.TestEnum.value_2)
+
+    def test_convert_value_from_member_names(self):
+        self.assertEqual(self.field.convert_value('value_1'), self.TestEnum.value_1)
+        self.assertEqual(self.field.convert_value('value_2'), self.TestEnum.value_2)
+
+    def test_export_definition(self):
+        self.assertEqual(self.field.export_definition(),
+                         {'alias': [],
+                          'doc': 'EnumField field (:class:`{0}`)'.format('.'.join([self.TestEnum.__module__,
+                                                                                   self.TestEnum.__name__])),
+                          'enum_class': self.TestEnum,
+                          'name': 'test_field', 'read_only': False},
+                         self.field.export_definition())