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