Skip to content

Instantly share code, notes, and snippets.

@mslabko
Forked from GrAndSE/enum.py
Created April 7, 2011 08:17
Show Gist options
  • Save mslabko/907321 to your computer and use it in GitHub Desktop.
Save mslabko/907321 to your computer and use it in GitHub Desktop.
ChoiceEnum is useful class to work with choices field ​​in Django. See http://docs.djangoproject.com/en/1.3/ref/models/fields/#choices
class Enum(object):
"""Unmodifiable enumeration class supports both sequential and named
creation:
>>> e = Enum('ZERO', 'ONE', TWO=2, FIVE=5)
>>> e.ZERO
0
>>> e.ONE
1
>>> e.FIVE
5
>>> tuple(e)
(('ZERO', 0), ('ONE', 1), ('TWO', 2), ('FIVE', 5))
"""
def __init__(self, *keys, **named):
"""Create new Enum instance
"""
super(Enum, self).__init__()
self.__dict__.update(zip(keys, range(len(keys))))
self.__dict__.update(named)
def __str__(self):
"""Get string representation
>>> e = Enum('ZERO', 'ONE', FIVE=5)
>>> print e
{
ZERO: 0,
ONE: 1,
FIVE: 5
}
"""
return "{\n%s\n}" % ",\n".join("%s: %s" % (key, self.__dict__[key])
for key in self.__dict__)
def __setattr__(self, name, value):
"""Enum is unmodifiable collection
>>> e = Enum('ZERO', 'ONE')
>>> e.FIVE = 6
Traceback (most recent call last):
File "<stdin>", line 1, in <module>
File "<stdin>", line 8, in __setattr__
TypeError: You can not modify enumerations
"""
raise TypeError('You can not modify enumerations')
def __delattr__(self, name, value):
"""Enum is unmodifiable collection
>>> e = Enum('ZERO', 'ONE')
>>> del e.ZERO
Traceback (most recent call last):
File "<stdin>", line 1, in <module>
File "<stdin>", line 8, in __setattr__
TypeError: You can not modify enumerations
"""
raise TypeError('You can not modify enumerations')
def __iter__(self):
"""Enumeration is iterable as a list of tuples (name, value)
>>> e = Enum('ZERO', 'ONE')
>>> tuple(e)
((ZERO, 0), (ONE, 1))
"""
return self.__dict__.iteritems()
class ChoiceEnum(Enum):
'''
Unmodifiable enumeration class for choices in DJANGO
Examle:
#if as_label=True (defalt) not tuple values of arguments represents as label
>>>COLORS = ChoiceEnum(red = ('r', 'red color'), blue=('blue color',), green = 'green color' )
>>>tuple(COLORS)
(('blue', 'blue color'), ('green', 'green color'), ('r', 'red color'))
#if as_label=False not tuple values of arguments represents as value
>>>COLORS = ChoiceEnum(red = ('r', 'red color'), blue=('b',), green = 'g' )
>>>tuple(COLORS)
(('b', 'blue'), ('g', 'green'), ('r', 'red color'))
>>>COLORS.red
'r'
>>>COLORS.red__label
'red color'
>>>COLORS['red']
'r'
>>COLORS.get_in('red', 'blue')
tuple('r', 'blue')
'''
def __init__(self, as_label=True, **named):
"""Create new ChoiceEnum instance
"""
super(ChoiceEnum, self).__init__()
self.__items = {}
d = {}
for i in named:
if type(named[i]) is not tuple:
d[i] = (i, named[i]) if as_label is True else (named[i], i)
elif len(named[i]) == 2:
d[i] = named[i]
elif len(named[i]) == 1:
d[i] = (i, named[i][0]) if as_label is True else (named[i][0], i)
else:
raise AttributeError('Invalid initial data of attributes')
named[i] = d[i][0]
self.__dict__.update( named )
self.__items.update( d )
def __setattr__(self, name, value):
"""Enum is unmodifiable collection
"""
if name != '_ChoiceEnum__items':
raise TypeError('You can not modify enumerations')
else:
return dict.__setattr__(self, name, value)
def __getitem__(self, name):
return self.__dict__[name][0]
def __getattribute__(self, name):
"""Get label for attribute
"""
try:
return object.__getattribute__(self, name)
except AttributeError:
if '__label' in name:
attr = name.split('__')
attr = attr[0]
return self.__items[attr][1]
pass
def label(self, name):
return self.__items[name][1] if name in self.__items else ''
def __iter__(self):
"""Enumeration is iterable as a list of tuples (name, value)
>>> e = Enum('ZERO', 'ONE')
>>> tuple(e)
((ZERO, 0), (ONE, 1))
"""
for key in self.__items:
yield (self.__items[key][0], self.__items[key][1])
def get_in(self, *attrs):
'''
return values of attributes in tuple
'''
return (self.__items[a][0] for a in attrs if a in self.__items)
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment