1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135 | #!/usr/bin/env python
class Choices(object):
"""
Convenience class for Django model choices which will use a [Small]IntegerField
for storage.
>>> statuses = Choices(
... ('live', 'Live'),
... ('draft', 'Draft'),
... ('hidden', 'Not Live'),
... )
>>> statuses.live
0
>>> statuses.hidden
2
>>> statuses.get_choices()
((0, 'Live'), (1, 'Draft'), (2, 'Not Live'))
This is then useful for use in a model field with the choices attribute.
>>> from django.db import models
>>> class Entry(models.Model):
... STATUSES = Choices(
... ('live', 'Live'),
... ('draft', 'Draft'),
... ('hidden', 'Not Live'),
... )
... status = models.SmallIntegerField(choices=STATUSES.get_choices(),
... default=STATUSES.live)
It's also useful later when you need to filter by your choices.
>>> live_entries = Entry.objects.filter(status=Entries.STATUSES.live)
"""
def __init__(self, *args):
super(Choices, self).__init__()
self.__dict__['_keys'], self.__dict__['_values'] = zip(*args)
def __getattr__(self, name):
# have to iterate manually to avoid conversion to a list for < 2.6 compat
for i, k in enumerate(self._keys):
if k == name:
return i
raise AttributeError("No attribute %r." % name)
def __setattr__(self, name, value):
raise AttributeError("%r object does not support attribute assignment" % self.__class__.__name__)
def __delattr__(self, name):
raise AttributeError("%r object does not support attribute deletion" % self.__class__.__name__)
def __getitem__(self, name):
try:
return self._values[getattr(self, name)]
except AttributeError:
raise KeyError(name)
def __contains__(self, name):
return name in self._keys
def __repr__(self):
return repr(self.items())
def items(self):
return tuple(zip(self._keys, self._values))
def keys(self):
# no need to copy since _keys is a tuple
return self._keys
def values(self):
# no need to copy since _values is a tuple
return self._values
def get_choices(self):
return tuple(enumerate(self._values))
# tests
if __name__ == '__main__':
import unittest
class ChoicesTests(unittest.TestCase):
def setUp(self):
self.c = Choices(
('stuff', 'Stuff'),
('whatnot', 'What-not'),
('things', 'Awesome Things'),
)
def test_choice_attributes(self):
self.assertEqual(self.c.stuff, 0)
self.assertEqual(self.c.whatnot, 1)
self.assertEqual(self.c.things, 2)
def test_choice_values(self):
self.assertEqual(self.c['stuff'], 'Stuff')
self.assertEqual(self.c['whatnot'], 'What-not')
self.assertEqual(self.c['things'], 'Awesome Things')
def test_model_choices(self):
correct_choices = (
(0, 'Stuff'),
(1, 'What-not'),
(2, 'Awesome Things'),
)
self.assertEqual(self.c.get_choices(), correct_choices)
def test_missing_attrs(self):
self.assertRaises(AttributeError, lambda: self.c.missing)
self.assertRaises(KeyError, lambda: self.c['missing'])
def _del_attr(self):
del self.c.stuff
def _del_item(self):
del self.c['stuff']
def _set_attr(self):
self.c.stuff = 'other stuff'
def _set_item(self):
self.c['stuff'] = 'other stuff'
def test_immutable(self):
self.assertRaises(AttributeError, self._del_attr)
self.assertRaises(AttributeError, self._set_attr)
self.assertRaises(TypeError, self._del_item)
self.assertRaises(TypeError, self._set_item)
unittest.main()
|
Comments
I updated the implementation to avoid the dependency on Django and to make the resulting instances immutable.
#
It should be nice to support grouped choices. Thus we'd do something like that:
#
Sorry I've made a mistake above:
Instead of:
#