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>

  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
-         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)
  
Allow string arguments (as slugs) when saving/updating EnumFields
## Code Before:
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)

## Instruction:
Allow string arguments (as slugs) when saving/updating EnumFields
## Code After:
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)
