You can not select more than 25 topics
Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.
85 lines
2.2 KiB
85 lines
2.2 KiB
7 months ago
|
# Lint as: python2, python3
|
||
|
# 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.
|
||
|
|
||
|
|
||
|
from __future__ import absolute_import
|
||
|
from __future__ import division
|
||
|
from __future__ import print_function
|
||
|
|
||
|
import six
|
||
|
from six.moves import zip
|
||
|
|
||
|
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 list(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, six.integer_types) 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
|