Enumeration field

  1
  2
  3
  4
  5
  6
  7
  8
  9
 10
 11
 12
 13
 14
 15
 16
 17
 18
 19
 20
 21
 22
 23
 24
 25
 26
 27
 28
 29
 30
 31
 32
 33
 34
 35
 36
 37
 38
 39
 40
 41
 42
 43
 44
 45
 46
 47
 48
 49
 50
 51
 52
 53
 54
 55
 56
 57
 58
 59
 60
 61
 62
 63
 64
 65
 66
 67
 68
 69
 70
 71
 72
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
# -*- coding: utf-8 -*-

from django.db import models

class Item(object):
    def __init__(self, value, slug, display=None):
        if not isinstance(value, int):
            raise TypeError('item value should be an integer, not %s' % value.__class__.__name__)
        if not isinstance(slug, str):
            raise TypeError('item slug should be a string, not %s' % slug.__class__.__name__)
        if display != None and not isinstance(display, (str)):
            raise TypeError('item slug should be a string, not %s' % display.__class__.__name__)
        super(Item, self).__init__()
        self.value = value
        self.slug = slug
        if display == None:
            self.display = slug
        else:
            self.display = display
        
    def __str__(self):
        return self.display
    
    def __repr__(self):
        return '<enum.Item: %s>' % self.display
    
    def __eq__(self, other):
        if isinstance(other, Item):
            return self.value == other.value        
        if isinstance(other, (int, str, unicode)):
            try:
                return self.value == int(other)
            except ValueError:
                return False        
        return False

    def __ne__(self, other):
        return not self.__eq__(other)


class Enumeration(object):
    @classmethod
    def from_value(cls, value):
        for attr in cls.__dict__.values():
            if isinstance(attr, Item) and attr.value == value:
                return attr
    
    @classmethod
    def from_slug(cls, slug):
        for attr in cls.__dict__.values():
            if isinstance(attr, Item) and attr.slug == slug:
                return attr
    
    @classmethod
    def get_items(cls):
        items = filter(lambda attr: isinstance(attr, Item), cls.__dict__.values())
        items.sort(lambda x, y: cmp(x.value, y.value))
        return items

    @classmethod
    def get_choices(cls):
        return [(item.value, item.display.capitalize()) for item in cls.get_items()]


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

    def __init__(self, enumeration, *args, **kwargs):
        kwargs.setdefault('choices', enumeration.get_choices())
        super(EnumField, self).__init__(*args, **kwargs)
        self.enumeration = enumeration
    
    def get_internal_type(self):
        return 'IntegerField'

    def to_python(self, value):
        if value == None or value == '' or value == u'':
            return None
        if isinstance(value, Item):
            return value
        if isinstance(value, int) or isinstance(value, str) or isinstance(value, unicode):
            item = self.enumeration.from_value(int(value))
            if item:
                return item
        
        raise ValueError, '%s is not a valid value for the enum field' % value
    
    
    def get_db_prep_save(self, value):
        if value:
            return value.value
    
    def get_db_prep_lookup(self, lookup_type, value):
        def prepare(value):
            if value == None:
                return None
            if isinstance(value, (int, str, unicode)):
                try:
                    return int(value)
                except ValueError:
                    raise ValueError('invalid value for the enum field lookup: %r' % value)
            if isinstance(value, Item):
                return value.value
        
        if lookup_type == 'exact':
            return [prepare(value)]
        elif lookup_type == 'in':
            return [prepare(v) for v in value]
        elif lookup_type == 'isnull':
            return []
        else:
            raise TypeError('Lookup type %r not supported.' % lookup_type)

#################
#
# tests/models.py
#
#################

# -*- coding: utf-8 -*-

from django.db import models
from utils import enum

class Animal(enum.Enumeration):
    CAT = enum.Item(2, 'cat', 'a cat')
    DOG = enum.Item(5, 'dog', 'a dog')

class MyModel(models.Model):
    animal = enum.EnumField(Animal)
    
    def __repr__(self):
        return '<MyModel: %s>' % self.animal


__test__ = {'API_TESTS': """
>>> model = MyModel()
>>> print model.animal
None

# The most straightforward way to assign a value is to use Item object
>>> model.animal = Animal.DOG
>>> model.animal
<enum.Item: a dog>

# But numeric value is also suitable
>>> model.animal = 2
>>> model.animal
<enum.Item: a cat>

# Even string representation of numeric value is suitable
>>> model.animal = '5'
>>> model.animal
<enum.Item: a dog>

# But if there no such numeric value, an exception will be raised
>>> model.animal = 28
Traceback (most recent call last):
    ...
ValueError: 28 is not a valid value for the enum field

# Empty string literal is valid and is equivalent to None value
>>> model.animal = ''
>>> print model.animal
None

# You can access to individual fields of an enum value easily
>>> model.animal = 2  #  its a cat
>>> model.animal.value
2
>>> model.animal.slug  # useful to use within URLs
'cat'
>>> model.animal.display  # useful for human-friendly display
'a cat'


# Lets test enum after save/restore from DB
>>> model.animal = Animal.CAT
>>> model.save()
>>> MyModel.objects.count()
1
>>> model = MyModel.objects.all()[0]
>>> model.animal
<enum.Item: a cat>

# Test lookup routines
>>> MyModel.objects.filter(animal=Animal.CAT)
[<MyModel: a cat>]
>>> MyModel.objects.filter(animal=2)
[<MyModel: a cat>]
>>> MyModel.objects.filter(animal='2')
[<MyModel: a cat>]
>>> MyModel.objects.filter(animal=28)
[]
>>> MyModel.objects.filter(animal='foo')
Traceback (most recent call last):
    ...
ValueError: invalid value for the enum field lookup: 'foo'

# 'in' lookup also works
>>> MyModel.objects.filter(animal__in=(Animal.CAT, Animal.DOG))
[<MyModel: a cat>]
>>> MyModel.objects.filter(animal__in=(Animal.DOG, 28, '34'))
[]
>>> MyModel.objects.filter(animal__in=(1, 2, 3))
[<MyModel: a cat>]

#########################
#   EnumField tests
#########################

# EnumField automatically creates choices. It uses capitalized display name
# and sorts by value in ascending order
>>> field = enum.EnumField(Animal)
>>> field.choices
[(2, 'A cat'), (5, 'A dog')]

#########################
#   Item tests
#########################

# Item constructor takes at least 2 arguments: integer 'value' and string 'slug'
>>> item = enum.Item(3, 'elephant')
>>> item.value
3
>>> item.slug
'elephant'

>>> item = enum.Item('3', 'elephant')
Traceback (most recent call last):
    ...
TypeError: item value should be an integer, not str

>>> item = enum.Item(3, 29010)
Traceback (most recent call last):
    ...
TypeError: item slug should be a string, not int


# Item takes optional third argument: 'display'. It should be a string. If display is ommited
# it takes value of slug
>>> item = enum.Item(3, 'elephant', 'a big elephant')
>>> item.display
'a big elephant'
>>> item = enum.Item(3, 'elephant')
>>> item.display
'elephant'

>>> item = enum.Item(3, 'elephant', 29010)
Traceback (most recent call last):
    ...
TypeError: item slug should be a string, not int

"""}

More like this

  1. Django enumeration for model field choices by martinthenext 2 years ago
  2. More readable Enumeration class for Django choices by achimnol 4 years, 9 months ago
  3. Choices helper by jacobian 4 years, 8 months ago
  4. Getting dynamic model choices in newforms by ubernostrum 7 years, 1 month ago
  5. Model Choices Helper by pmclanahan 4 years, 2 months ago

Comments

statico (on March 6, 2009):

I had to modify this slightly to work with admin screens correctly, but otherwise this is fantastic. Thank you!

To make this work in the admin listings, to_python() needed to return a key in the EnumField's choices property -- the enum's integer value -- instead of the Item object:

def to_python(self, value):
    if value == None or value == '' or value == u'':
        return None
    if isinstance(value, Item):
        return value.value  # Changed
    if isinstance(value, int) or isinstance(value, str) or isinstance(value, unicode):
        item = self.enumeration.from_value(int(value))
        if item:
            return item.value  # Changed

    raise ValueError, '%s is not a valid value for the enum field' % value

The dereferencing of the enum's int value is no longer needed, thus:

def get_db_prep_save(self, value):
    if value:
        return value  # Changed

Unfortunately, this means I'll have to use Enumeration.from_value() more often, but I'm okay with that.

Thinking about it now, I wonder if get_choices() could return a tuple of (item, item.display) instead of (item.value, item.display)...

#

gpothier (on September 20, 2010):

I can confirm that having get_choices() return a tuple of (item, item.display) also permits to have the admin pages working, with the added benefit of not having to call Enumeration.from_value() more often. The only other change is to make Item.str() return self.value instead of self.display.

Of course this is not a perfect solution either as directly printing an Item will just print its value.

#

gpothier (on September 20, 2010):

Well, there is a better solution: implement Item.unicode() and have it return the value, and keep Item.str() returning the display. It works in forms, and prints the display name when directly printing the Item.

#

demichej (on August 3, 2012):

Thanks, gpothier. That worked for me. I'm curious why the original author hasn't updated the snippet?

#

demichej (on August 3, 2012):

Hmmm. Yeah, this actually doesn't work at all for Django 1.4. I'm working on tweaking my version, but if someone else has this working for 1.4, please let me know. Thanks.

#

mmichele (on March 29, 2014):

Thanks for your code, is very useful. But there is an error on parameters of get_db_prep_save function with django 1.5.

Please change 89 line from

def get_db_prep_save(self, value):

to

def get_db_prep_save(self, value, connection):

Thanks!

#

(Forgotten your password?)