# -*- 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 '' % 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 '' % 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 # But numeric value is also suitable >>> model.animal = 2 >>> model.animal # Even string representation of numeric value is suitable >>> model.animal = '5' >>> model.animal # 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 # Test lookup routines >>> MyModel.objects.filter(animal=Animal.CAT) [] >>> MyModel.objects.filter(animal=2) [] >>> MyModel.objects.filter(animal='2') [] >>> 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.objects.filter(animal__in=(Animal.DOG, 28, '34')) [] >>> MyModel.objects.filter(animal__in=(1, 2, 3)) [] ######################### # 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 """}