choices に enum を指定可能なカスタムフィールドを作る

python 3.4 で 導入予定、かつ python 2.4 にバックポートされた enum34 が便利なので CharField の choices オプションに指定しようとしたけど、そのままじゃ上手く使えなかったので作ってみた。

参考URL

- http://qiita.com/methane/items/8612bdefd8fa4238cc44
- http://docs.djangoproject.jp/en/latest/ref/models/fields.html

動作環境

- Python 2.7.4
- Django 1.6.2

```py:choicefield.py
# -*- encoding: utf-8 -*-
from django.db import models
from django.utils.functional import curry
import enum


class ChoiceField(models.CharField):
    u""" choices に enum.Enum のサブクラスを指定できるカスタムフィールドです。
    """
    __metaclass__ = models.SubfieldBase

    def __init__(self, *args, **kwargs):
        self.enum = kwargs.get("choices", None)
        if self.enum and issubclass(self.enum, enum.Enum):
            kwargs["choices"] = self._from_enum()
            kwargs["max_length"] = self._calc_max_length()
        super(ChoiceField, self).__init__(*args, **kwargs)

    def _calc_max_length(self):
        u""" DBに保存するのに必要な文字列長を求める。
        """
        return max([len(unicode(item)) for item in self.enum])

    def _from_enum(self):
        u""" choices に指定できるタプルに変換する。
        """
        return [(item, item.value) for item in self.enum]

    @staticmethod
    def _get_display(self, field):
        u""" 表示用にモデルに後付けするメソッド
        """
        return getattr(self, field.attname).value

    def contribute_to_class(self, cls, name, virtual_only=False):
        super(ChoiceField, self).contribute_to_class(cls, name, virtual_only)
        setattr(cls, 'get_%s_display' % self.name, curry(self._get_display, field=self))
        # MaxLengthValidator で必要になるため Enum クラスに後付けする
        setattr(self.enum, '__len__', lambda x: len(unicode(x)))

    def get_prep_value(self, value):
        if isinstance(value, basestring):
            return value
        return None if value is None else unicode(value)

    def to_python(self, value):
        if not value or not isinstance(value, basestring):
            return value

        try:
            return self.enum[value]
        except KeyError:
            for m in self.enum:
                if value == m:
                    return value
                if value == m.value:
                    return m
                if value == m.name:
                    return m
                if value.endswith(m.name):
                    return m
            raise Exception('%s is not a valid value for enum %s' % (value, self.enum))
```

テストはこんな感じ

```py:tests/tests_choicefield.py
# -*- encoding: utf-8 -*-
from django import test
from django.db import models
import enum
import os


class Gender(enum.Enum):
    __order__ = "Male Female"
    Male = u"男性"
    Female = u"女性"


class Person(models.Model):
    name = models.CharField(u"名前", max_length=255)
    gender = choices.ChoiceField(u"性別", choices=Gender, max_length=255)

    class Meta:
        app_label = os.path.basename(os.path.abspath(os.path.join(os.path.split(__file__)[0], os.pardir)))


class PersonTest(test.TestCase):
    def test_choices(self):
        obj = Person.objects.create(name="John", gender=Gender.Male)
        self.assertEqual(Gender.Male, obj.gender)
        self.assertEqual(u"男性", obj.get_gender_display())

    def test_get(self):
        obj1 = Person.objects.create(name="John", gender=Gender.Male)
        obj2 = Person.objects.get(name="John", gender=Gender.Male)
        self.assertEqual(obj1, obj2)

    def test_filter_in(self):
        obj1 = Person.objects.create(name="John", gender=Gender.Male)
        obj2 = Person.objects.create(name="Jane", gender=Gender.Female)
        actual = Person.objects.filter(gender__in=[Gender.Male, Gender.Female])
        self.assertItemsEqual([obj1, obj2], actual)
```