1# Lint as: python2, python3 2# Copyright 2020 The Chromium OS Authors. All rights reserved. 3# Use of this source code is governed by a BSD-style license that can be 4# found in the LICENSE file. 5 6 7from __future__ import absolute_import 8from __future__ import division 9from __future__ import print_function 10 11import six 12from six.moves import zip 13 14class AutotestEnum(object): 15 """ 16 Utility class to implement Enum-like functionality. 17 18 >>> e = Enum('String one', 'String two') 19 >>> e.STRING_ONE 20 0 21 >>> e.STRING_TWO 22 1 23 >>> e.choices() 24 [(0, 'String one'), (1, 'String two')] 25 >>> e.get_value('String one') 26 0 27 >>> e.get_string(0) 28 'String one' 29 30 >>> e = Enum('Hello', 'Goodbye', string_values=True) 31 >>> e.HELLO, e.GOODBYE 32 ('Hello', 'Goodbye') 33 34 >>> e = Enum('One', 'Two', start_value=1) 35 >>> e.ONE 36 1 37 >>> e.TWO 38 2 39 """ 40 41 def __init__(self, *names, **kwargs): 42 self.string_values = kwargs.get('string_values') 43 start_value = kwargs.get('start_value', 0) 44 step = kwargs.get('step', 1) 45 self.names = names 46 self.values = [] 47 for i, name in enumerate(names): 48 if self.string_values: 49 value = name 50 else: 51 value = i * step + start_value 52 self.values.append(value) 53 setattr(self, self.get_attr_name(name), value) 54 55 56 @staticmethod 57 def get_attr_name(string): 58 return string.upper().replace(' ', '_') 59 60 61 def choices(self): 62 'Return choice list suitable for Django model choices.' 63 return list(zip(self.values, self.names)) 64 65 66 def get_value(self, name): 67 """\ 68 Convert a string name to it's corresponding value. If a value 69 is passed in, it is returned. 70 """ 71 if isinstance(name, six.integer_types) and not self.string_values: 72 # name is already a value 73 return name 74 return getattr(self, self.get_attr_name(name)) 75 76 77 def get_string(self, value): 78 ' Given a value, get the string name for it.' 79 if value not in self.values: 80 raise ValueError('Value %s not in this enum' % value) 81 index = self.values.index(value) 82 return self.names[index] 83 84Enum = AutotestEnum 85