Merge pull request #74 from akx/drf-field

Django REST Framework support

Author: akx

enumfields/drf/__init__.py

@@ -0,0 +1,2 @@
+from .fields import EnumField
+from .serializers import EnumSupportSerializerMixin

enumfields/drf/fields.py

@@ -0,0 +1,57 @@
+from django.utils import six
+from django.utils.encoding import force_text
+from rest_framework.fields import ChoiceField
+
+
+class EnumField(ChoiceField):
+    def __init__(self, enum, lenient=False, ints_as_names=False, **kwargs):
+        """
+        :param enum: The enumeration class. 
+        :param lenient: Whether to allow lenient parsing (case-insensitive, by value or name)
+        :type lenient: bool
+        :param ints_as_names: Whether to serialize integer-valued enums by their name, not the integer value
+        :type ints_as_names: bool
+        """
+        self.enum = enum
+        self.lenient = lenient
+        self.ints_as_names = ints_as_names
+        kwargs['choices'] = tuple((e.value, getattr(e, 'label', e.name)) for e in self.enum)
+        super(EnumField, self).__init__(**kwargs)
+
+    def to_representation(self, instance):
+        if instance in ('', u'', None):
+            return instance
+        try:
+            if not isinstance(instance, self.enum):
+                instance = self.enum(instance)  # Try to cast it
+            if self.ints_as_names and isinstance(instance.value, six.integer_types):
+                # If the enum value is an int, assume the name is more representative
+                return instance.name.lower()
+            return instance.value
+        except ValueError:
+            raise ValueError('Invalid value [%r] of enum %s' % (instance, self.enum.__name__))
+
+    def to_internal_value(self, data):
+        if isinstance(data, self.enum):
+            return data
+        try:
+            # Convert the value using the same mechanism DRF uses
+            converted_value = self.choice_strings_to_values[six.text_type(data)]
+            return self.enum(converted_value)
+        except (ValueError, KeyError):
+            pass
+
+        if self.lenient:
+            # Normal logic:
+            for choice in self.enum:
+                if choice.name == data or choice.value == data:
+                    return choice
+
+            # Case-insensitive logic:
+            l_data = force_text(data).lower()
+            for choice in self.enum:
+                if choice.name.lower() == l_data or force_text(choice.value).lower() == l_data:
+                    return choice
+
+        # Fallback (will likely just raise):
+        return super(EnumField, self).to_internal_value(data)

enumfields/drf/serializers.py

@@ -0,0 +1,18 @@
+from rest_framework.fields import ChoiceField
+
+from enumfields.drf.fields import EnumField as EnumSerializerField
+from enumfields.fields import EnumFieldMixin
+
+
+class EnumSupportSerializerMixin(object):
+    enumfield_options = {}
+
+    def build_standard_field(self, field_name, model_field):
+        field_class, field_kwargs = (
+            super(EnumSupportSerializerMixin, self).build_standard_field(field_name, model_field)
+        )
+        if field_class == ChoiceField and isinstance(model_field, EnumFieldMixin):
+            field_class = EnumSerializerField
+            field_kwargs['enum'] = model_field.enum
+            field_kwargs.update(self.enumfield_options)
+        return field_class, field_kwargs

setup.py

@@ -55,6 +55,7 @@ def run_tests(self):
     tests_require=[
         'pytest-django',
         'Django',
+        'djangorestframework'
     ],
     extras_require={
         ":python_version<'3.4'": ['enum34'],

tests/test_serializers.py

@@ -0,0 +1,77 @@
+# -- encoding: UTF-8 --
+
+import uuid
+
+import pytest
+from rest_framework import serializers
+
+from enumfields.drf.serializers import EnumSupportSerializerMixin
+
+from .enums import Color, IntegerEnum, Taste
+from .models import MyModel
+
+
+class MySerializer(EnumSupportSerializerMixin, serializers.ModelSerializer):
+    class Meta:
+        model = MyModel
+        fields = '__all__'
+
+
+class LenientIntNameSerializer(MySerializer):
+    enumfield_options = {
+        'lenient': True,
+        'ints_as_names': True,
+    }
+
+
+@pytest.mark.parametrize('int_names', (False, True))
+def test_serialize(int_names):
+    inst = MyModel(color=Color.BLUE, taste=Taste.UMAMI, int_enum=IntegerEnum.B)
+    data = (LenientIntNameSerializer if int_names else MySerializer)(inst).data
+    assert data['color'] == Color.BLUE.value
+    if int_names:
+        assert data['taste'] == 'umami'
+        assert data['int_enum'] == 'b'
+    else:
+        assert data['taste'] == Taste.UMAMI.value
+        assert data['int_enum'] == IntegerEnum.B.value
+
+
+@pytest.mark.django_db
+@pytest.mark.parametrize('lenient_serializer', (False, True))
+@pytest.mark.parametrize('lenient_data', (False, True))
+def test_deserialize(lenient_data, lenient_serializer):
+    secret_uuid = str(uuid.uuid4())
+    data = {
+        'color': Color.BLUE.value,
+        'taste': Taste.UMAMI.value,
+        'int_enum': IntegerEnum.B.value,
+        'random_code': secret_uuid,
+    }
+    if lenient_data:
+        data.update({
+            'color': 'b',
+            'taste': 'Umami',
+            'int_enum': 'B',
+        })
+    serializer_cls = (LenientIntNameSerializer if lenient_serializer else MySerializer)
+    serializer = serializer_cls(data=data)
+    if lenient_data and not lenient_serializer:
+        assert not serializer.is_valid()
+        return
+    assert serializer.is_valid(), serializer.errors
+
+    validated_data = serializer.validated_data
+    assert validated_data['color'] == Color.BLUE
+    assert validated_data['taste'] == Taste.UMAMI
+    assert validated_data['int_enum'] == IntegerEnum.B
+
+    inst = serializer.save()
+    assert inst.color == Color.BLUE
+    assert inst.taste == Taste.UMAMI
+    assert inst.int_enum == IntegerEnum.B
+
+    inst = MyModel.objects.get(random_code=secret_uuid)  # will raise if fails
+    assert inst.color == Color.BLUE
+    assert inst.taste == Taste.UMAMI
+    assert inst.int_enum == IntegerEnum.B