blob: 601044afe8db3ee26735f7575b9b613c303d6619 [file] [log] [blame]
# Copyright 2020 The Chromium OS Authors. All rights reserved.
# Use of this source code is governed by a BSD-style license that can be
# found in the LICENSE file.
class AutotestEnum(object):
"""
Utility class to implement Enum-like functionality.
>>> e = Enum('String one', 'String two')
>>> e.STRING_ONE
0
>>> e.STRING_TWO
1
>>> e.choices()
[(0, 'String one'), (1, 'String two')]
>>> e.get_value('String one')
0
>>> e.get_string(0)
'String one'
>>> e = Enum('Hello', 'Goodbye', string_values=True)
>>> e.HELLO, e.GOODBYE
('Hello', 'Goodbye')
>>> e = Enum('One', 'Two', start_value=1)
>>> e.ONE
1
>>> e.TWO
2
"""
def __init__(self, *names, **kwargs):
self.string_values = kwargs.get('string_values')
start_value = kwargs.get('start_value', 0)
step = kwargs.get('step', 1)
self.names = names
self.values = []
for i, name in enumerate(names):
if self.string_values:
value = name
else:
value = i * step + start_value
self.values.append(value)
setattr(self, self.get_attr_name(name), value)
@staticmethod
def get_attr_name(string):
return string.upper().replace(' ', '_')
def choices(self):
'Return choice list suitable for Django model choices.'
return zip(self.values, self.names)
def get_value(self, name):
"""\
Convert a string name to it's corresponding value. If a value
is passed in, it is returned.
"""
if isinstance(name, (int, long)) and not self.string_values:
# name is already a value
return name
return getattr(self, self.get_attr_name(name))
def get_string(self, value):
' Given a value, get the string name for it.'
if value not in self.values:
raise ValueError('Value %s not in this enum' % value)
index = self.values.index(value)
return self.names[index]
Enum = AutotestEnum