1 """\ 2 Generic enumeration support. 3 """ 4 5 __author__ = 'showard (at] google.com (Steve Howard)' 6 7 class 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
