• Home
  • Line#
  • Scopes#
  • Navigate#
  • Raw
  • Download
1"""\
2Generic enumeration support.
3"""
4
5__author__ = 'showard@google.com (Steve Howard)'
6
7class Enum(object):
8    """\
9    Utility class to implement Enum-like functionality.
10
11    >>> e = Enum('String one', 'String two')
12    >>> e.STRING_ONE
13    0
14    >>> e.STRING_TWO
15    1
16    >>> e.choices()
17    [(0, 'String one'), (1, 'String two')]
18    >>> e.get_value('String one')
19    0
20    >>> e.get_string(0)
21    'String one'
22
23    >>> e = Enum('Hello', 'Goodbye', string_values=True)
24    >>> e.HELLO, e.GOODBYE
25    ('Hello', 'Goodbye')
26
27    >>> e = Enum('One', 'Two', start_value=1)
28    >>> e.ONE
29    1
30    >>> e.TWO
31    2
32    """
33    def __init__(self, *names, **kwargs):
34        self.string_values = kwargs.get('string_values')
35        start_value = kwargs.get('start_value', 0)
36        step = kwargs.get('step', 1)
37        self.names = names
38        self.values = []
39        for i, name in enumerate(names):
40            if self.string_values:
41                value = name
42            else:
43                value = i * step + start_value
44            self.values.append(value)
45            setattr(self, self.get_attr_name(name), value)
46
47
48    @staticmethod
49    def get_attr_name(string):
50        return string.upper().replace(' ', '_')
51
52
53    def choices(self):
54        'Return choice list suitable for Django model choices.'
55        return zip(self.values, self.names)
56
57
58    def get_value(self, name):
59        """\
60        Convert a string name to it's corresponding value.  If a value
61        is passed in, it is returned.
62        """
63        if isinstance(name, (int, long)) and not self.string_values:
64            # name is already a value
65            return name
66        return getattr(self, self.get_attr_name(name))
67
68
69    def get_string(self, value):
70        ' Given a value, get the string name for it.'
71        if value not in self.values:
72            raise ValueError('Value %s not in this enum' % value)
73        index = self.values.index(value)
74        return self.names[index]
75