From a2dcac06deec9d360c06dae37c2e706e970a019a Mon Sep 17 00:00:00 2001
From: Carl Meyer <carl@dirtcircle.com>
Date: Tue, 12 Jan 2010 21:47:28 -0500
Subject: [PATCH] added ChoiceEnum

---
 README.txt                 | 29 ++++++++++++++++++++++++
 model_utils/__init__.py    | 45 ++++++++++++++++++++++++++++++++++++++
 model_utils/tests/tests.py | 14 ++++++++++++
 3 files changed, 88 insertions(+)

diff --git a/README.txt b/README.txt
index 18a34fe..7d5526c 100644
--- a/README.txt
+++ b/README.txt
@@ -28,6 +28,35 @@ Dependencies
 
 .. _Django: http://www.djangoproject.com/
 
+ChoiceEnum
+==========
+
+``ChoiceEnum`` makes setting ``choices`` on a Django model field way
+too easy::
+
+    from model_utils import ChoiceEnum
+
+    class Article(models.Model):
+        STATUS = ChoiceEnum('draft', 'published')
+        # ...
+        status = models.PositiveIntegerField(choices=STATUS, default=STATUS.draft)
+
+        def status_desc(self):
+            return self.STATUS[self.status]
+
+A ``ChoiceEnum`` object is initialized with any number of choices,
+which should be strings. It assigns a sequential id to each
+choice. The numerical id for a choice is available through attribute
+access (``STATUS.draft``), and the text name for a choice can be
+obtained by indexing with the numerical id
+(``self.STATUS[self.status]``). If iterated over, a ``ChoiceEnum``
+object yields a tuple of two-tuples linking id to text names, the
+format expected by the ``choices`` attribute of Django models.
+
+Be careful not to add new choices in the middle of the list, as that
+will change the numerical ids for all subsequent choices, which could
+impact existing data.
+
 models.InheritanceCastModel
 ===========================
 
diff --git a/model_utils/__init__.py b/model_utils/__init__.py
index e69de29..18991e7 100644
--- a/model_utils/__init__.py
+++ b/model_utils/__init__.py
@@ -0,0 +1,45 @@
+class ChoiceEnum(object):
+    """
+    A class to encapsulate handy functionality for lists of choices
+    for a Django model field.
+
+    Accepts verbose choice names as arguments, and automatically
+    assigns numeric keys to them. When iterated over, behaves as the
+    standard Django choices tuple of two-tuples.
+
+    Attribute access allows conversion of verbose choice name to
+    choice key, dictionary access the reverse.
+
+    Example:
+
+    >>> STATUS = ChoiceEnum('DRAFT', 'PUBLISHED')
+    >>> STATUS.DRAFT
+    0
+    >>> STATUS[1]
+    'PUBLISHED'
+    >>> tuple(STATUS)
+    ((0, 'DRAFT'), (1, 'PUBLISHED'))
+    
+    """
+    def __init__(self, *choices):
+        self._choices = tuple(enumerate(choices))
+        self._choice_dict = dict(self._choices)
+        self._reverse_dict = dict(((i[1], i[0]) for i in self._choices))
+
+    def __iter__(self):
+        return iter(self._choices)
+
+    def __getattr__(self, attname):
+        try:
+            return self._reverse_dict[attname]
+        except KeyError:
+            raise AttributeError(attname)
+
+    def __getitem__(self, key):
+        return self._choice_dict[key]
+
+    def __repr__(self):
+        return '%s(%s)' % (self.__class__.__name__,
+                           ', '.join(("'%s'" % i[1] for i in self._choices)))
+    
+
diff --git a/model_utils/tests/tests.py b/model_utils/tests/tests.py
index 40dcebd..197f7c3 100644
--- a/model_utils/tests/tests.py
+++ b/model_utils/tests/tests.py
@@ -1,9 +1,23 @@
 from django.test import TestCase
 from django.contrib.contenttypes.models import ContentType
 
+from model_utils import ChoiceEnum
 from model_utils.tests.models import InheritParent, InheritChild, TimeStamp, \
     Post
+    
+class ChoiceEnumTests(TestCase):
+    def setUp(self):
+        self.STATUS = ChoiceEnum('DRAFT', 'PUBLISHED')
 
+    def test_getattr(self):
+        self.assertEquals(self.STATUS.DRAFT, 0)
+
+    def test_getitem(self):
+        self.assertEquals(self.STATUS[1], 'PUBLISHED')
+
+    def test_iteration(self):
+        self.assertEquals(tuple(self.STATUS), ((0, 'DRAFT'), (1, 'PUBLISHED')))
+    
 class InheritanceCastModelTests(TestCase):
     def setUp(self):
         self.parent = InheritParent.objects.create()