Django snippets

# -*- 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

"""}