diff --git a/model_utils/__init__.py b/model_utils/__init__.py index 35b1a68..b9202ce 100644 --- a/model_utils/__init__.py +++ b/model_utils/__init__.py @@ -57,39 +57,70 @@ class Choices(object): A class to encapsulate handy functionality for lists of choices for a Django model field. - Accepts as arguments either tuples mapping choice IDs (strings) to - human-readable names, or simply choice IDs (in which case the ID - is also used as the human-readable name). When iterated over, + Each argument to ``Choices`` is a choice, represented as either a + string, a two-tuple, or a three-tuple. + + If a single string is provided, that string is used as the + database representation of the choice as well as the + human-readable presentation. + + If a two-tuple is provided, the first item is used as the database + representation and the second the human-readable presentation. + + If a triple is provided, the first item is the database + representation, the second a valid Python identifier that can be + used as a readable label in code, and the third the human-readable + presentation. This is most useful when the database representation + must sacrifice readability for some reason: to achieve a specific + ordering, to use an integer rather than a character field, etc. + + Regardless of what representation of each choice is originally + given, when iterated over or indexed into, a ``Choices`` object behaves as the standard Django choices list of two-tuples. - Choice IDs can be accessed as attributes for readable code. + If the triple form is used, the Python identifier names can be + accessed as attributes on the ``Choices`` object, returning the + database representation. (If the single or two-tuple forms are + used and the database representation happens to be a valid Python + identifier, the database representation itself is available as an + attribute on the ``Choices`` object, returning itself.) """ def __init__(self, *choices): - self._choices = list(self.equalize(choices)) - self._choice_dict = dict(self._choices) - self._reverse_dict = dict(((i[0], i[0]) for i in self._choices)) + self._full = [] + self._choices = [] + self._choice_dict = {} + for choice in self.equalize(choices): + self._full.append(choice) + self._choices.append((choice[0], choice[2])) + self._choice_dict[choice[1]] = choice[0] def equalize(self, choices): for choice in choices: if isinstance(choice, (list, tuple)): - yield choice + if len(choice) == 3: + yield choice + elif len(choice) == 2: + yield (choice[0], choice[0], choice[1]) + else: + raise ValueError("Choices can't handle a list/tuple of length %s, only 2 or 3" + % len(choice)) else: - yield (choice, choice) + yield (choice, choice, choice) def __iter__(self): return iter(self._choices) def __getattr__(self, attname): try: - return self._reverse_dict[attname] + return self._choice_dict[attname] except KeyError: - raise AttributeError(attname) - + raise AttributeError(attname) + def __getitem__(self, index): return self._choices[index] def __repr__(self): return '%s(%s)' % (self.__class__.__name__, - ', '.join(("'%s'" % i[0] for i in self._choices))) + ', '.join(("%s" % str(i) for i in self._full))) diff --git a/model_utils/tests/tests.py b/model_utils/tests/tests.py index 609f796..2b3fef9 100644 --- a/model_utils/tests/tests.py +++ b/model_utils/tests/tests.py @@ -111,6 +111,12 @@ class ChoicesTests(TestCase): def test_iteration(self): self.assertEquals(tuple(self.STATUS), (('DRAFT', 'DRAFT'), ('PUBLISHED', 'PUBLISHED'))) + def test_repr(self): + self.assertEquals(repr(self.STATUS), + "Choices(" + "('DRAFT', 'DRAFT', 'DRAFT'), " + "('PUBLISHED', 'PUBLISHED', 'PUBLISHED'))") + class LabelChoicesTests(ChoicesTests): def setUp(self): @@ -136,7 +142,41 @@ class LabelChoicesTests(ChoicesTests): def test_provided(self): self.assertEquals(self.STATUS.DRAFT, 'DRAFT') + def test_repr(self): + self.assertEquals(repr(self.STATUS), + "Choices(" + "('DRAFT', 'DRAFT', 'is draft'), " + "('PUBLISHED', 'PUBLISHED', 'is published'), " + "('DELETED', 'DELETED', 'DELETED'))") + +class IdentifierChoicesTests(ChoicesTests): + def setUp(self): + self.STATUS = Choices( + (0, 'DRAFT', 'is draft'), + (1, 'PUBLISHED', 'is published'), + (2, 'DELETED', 'is deleted')) + + def test_iteration(self): + self.assertEqual(tuple(self.STATUS), ( + (0, 'is draft'), + (1, 'is published'), + (2, 'is deleted'))) + + def test_indexing(self): + self.assertEquals(self.STATUS[1], (1, 'is published')) + + def test_getattr(self): + self.assertEquals(self.STATUS.DRAFT, 0) + + def test_repr(self): + self.assertEquals(repr(self.STATUS), + "Choices(" + "(0, 'DRAFT', 'is draft'), " + "(1, 'PUBLISHED', 'is published'), " + "(2, 'DELETED', 'is deleted'))") + + class InheritanceCastModelTests(TestCase): def setUp(self): self.parent = InheritParent.objects.create()