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