from django.utils.datastructures import SortedDict class Choices(object): """ Easy declarative "choices" tool:: >>> STATUSES = Choices("Live", "Draft") # Acts like a choices list: >>> list(STATUSES) [(1, 'Live'), (2, 'Draft')] # Easily convert from code to verbose: >>> STATUSES.verbose(1) 'Live' # ... and vice versa: >>> STATUSES.code("Draft") 2 """ def __init__(self, *args, **kwargs): self.code_map = SortedDict() self.verbose_map = {} for code, verbose in enumerate(args): # Enumerate starts from 0, but for convention's sake we'd prefer to # start choices from 1. self.code_map[code+1] = verbose self.verbose_map[verbose] = code+1 for code, verbose in kwargs.items(): self.code_map[code] = verbose self.verbose_map[verbose] = code def __iter__(self): return self.code_map.iteritems() def __len__(self): return len(self.code_map) def code(self, verbose): """ Return the code version of the verbose name. """ return self.verbose_map[verbose] def verbose(self, code): """ Return the verbose name given the code. """ return self.code_map[code] # # Tests # import unittest class ChoicesTests(unittest.TestCase): def test_positional_choices(self): ch = Choices("Live", "Draft") self.assertEqual(list(ch), [(1, 'Live'), (2, 'Draft')]) self.assertEqual(ch.verbose(1), 'Live') self.assertEqual(ch.verbose(2), 'Draft') self.assertEqual(ch.code('Live'), 1) self.assertEqual(ch.code('Draft'), 2) def test_keyword_choices(self): ch = Choices(l="Live", d="Draft") self.assertEqual(sorted(list(ch)), [('d', 'Draft'), ('l', 'Live')]) self.assertEqual(ch.verbose('l'), 'Live') self.assertEqual(ch.verbose('d'), 'Draft') self.assertEqual(ch.code('Live'), 'l') self.assertEqual(ch.code('Draft'), 'd') if __name__ == '__main__': unittest.main()