from django.db import models

class EnumField(models.Field):
    __metaclass__ = models.SubfieldBase

    def __init__(self, enumeration, *args, **kwargs):
        self.enumeration = enumeration
        kwargs.setdefault('choices', enumeration.get_choices())

        super(EnumField, self).__init__(*args, **kwargs)

    def get_internal_type(self):
        return 'IntegerField'

    def to_python(self, value):
        return self.enumeration.to_item(value)

    def get_db_prep_save(self, value, connection=None):
        if hasattr(value, 'value'):
            return value.value
        return value

    def get_db_prep_lookup(self, lookup_type, value, connection=None, prepared=False):
        def prepare(value):
            v = self.to_python(value)
            return self.get_db_prep_save(v, connection=connection)

        if lookup_type == 'exact':
            return [prepare(value)]
        elif lookup_type == 'in':
            return [prepare(v) for v in value]
        elif lookup_type == 'isnull':
            return []

        raise TypeError("Lookup type %r not supported." % lookup_type)

from django.db import models

class EnumField(models.Field):
    __metaclass__ = models.SubfieldBase

    def __init__(self, enumeration, *args, **kwargs):
        self.enumeration = enumeration
        kwargs.setdefault('choices', enumeration.get_choices())

        super(EnumField, self).__init__(*args, **kwargs)

    def get_internal_type(self):
        return 'IntegerField'

    def to_python(self, value):
        return self.enumeration.to_item(value)

    def get_db_prep_save(self, value, connection=None):
        return self.to_python(value).value

    def get_db_prep_lookup(self, lookup_type, value, connection=None, prepared=False):
        def prepare(value):
            v = self.to_python(value)
            return self.get_db_prep_save(v, connection=connection)

        if lookup_type == 'exact':
            return [prepare(value)]
        elif lookup_type == 'in':
            return [prepare(v) for v in value]
        elif lookup_type == 'isnull':
            return []

        raise TypeError("Lookup type %r not supported." % lookup_type)

Allow string arguments (as slugs) when saving/updating EnumFields
Allow string arguments (as slugs) when saving/updating EnumFields

This fixes issues where:

  MyModel.objects.update(my_enum_field='slug')

would result in SQL like:

  UPDATE app_mymodel SET my_enum_field = 'slug'

.. instead of what that's slug's integer value is.

Signed-off-by: Chris Lamb <711c73f64afdce07b7e38039a96d2224209e9a6c@playfire.com>
