1# Copyright 2011-2019, Damian Johnson and The Tor Project
2# See LICENSE for licensing information
3
4"""
5Basic enumeration, providing ordered types for collections. These can be
6constructed as simple type listings...
7
8::
9
10  >>> from stem.util import enum
11  >>> insects = enum.Enum('ANT', 'WASP', 'LADYBUG', 'FIREFLY')
12  >>> insects.ANT
13  'Ant'
14  >>> tuple(insects)
15  ('Ant', 'Wasp', 'Ladybug', 'Firefly')
16
17... or with overwritten string counterparts...
18
19::
20
21  >>> from stem.util import enum
22  >>> pets = enum.Enum(('DOG', 'Skippy'), 'CAT', ('FISH', 'Nemo'))
23  >>> pets.DOG
24  'Skippy'
25  >>> pets.CAT
26  'Cat'
27
28**Module Overview:**
29
30::
31
32  UppercaseEnum - Provides an enum instance with capitalized values
33
34  Enum - Provides a basic, ordered  enumeration
35    |- keys - string representation of our enum keys
36    |- index_of - index of an enum value
37    |- next - provides the enum after a given enum value
38    |- previous - provides the enum before a given value
39    |- __getitem__ - provides the value for an enum key
40    +- __iter__ - iterator over our enum keys
41"""
42
43import stem.util
44
45
46def UppercaseEnum(*args):
47  """
48  Provides an :class:`~stem.util.enum.Enum` instance where the values are
49  identical to the keys. Since the keys are uppercase by convention this means
50  the values are too. For instance...
51
52  ::
53
54    >>> from stem.util import enum
55    >>> runlevels = enum.UppercaseEnum('DEBUG', 'INFO', 'NOTICE', 'WARN', 'ERROR')
56    >>> runlevels.DEBUG
57    'DEBUG'
58
59  :param list args: enum keys to initialize with
60
61  :returns: :class:`~stem.util.enum.Enum` instance with the given keys
62  """
63
64  return Enum(*[(v, v) for v in args])
65
66
67class Enum(object):
68  """
69  Basic enumeration.
70  """
71
72  def __init__(self, *args):
73    from stem.util.str_tools import _to_camel_case
74
75    # ordered listings of our keys and values
76    keys, values = [], []
77
78    for entry in args:
79      if stem.util._is_str(entry):
80        key, val = entry, _to_camel_case(entry)
81      elif isinstance(entry, tuple) and len(entry) == 2:
82        key, val = entry
83      else:
84        raise ValueError('Unrecognized input: %s' % args)
85
86      keys.append(key)
87      values.append(val)
88      setattr(self, key, val)
89
90    self._keys = tuple(keys)
91    self._values = tuple(values)
92
93  def keys(self):
94    """
95    Provides an ordered listing of the enumeration keys in this set.
96
97    :returns: **list** with our enum keys
98    """
99
100    return list(self._keys)
101
102  def index_of(self, value):
103    """
104    Provides the index of the given value in the collection.
105
106    :param str value: entry to be looked up
107
108    :returns: **int** index of the given entry
109
110    :raises: **ValueError** if no such element exists
111    """
112
113    return self._values.index(value)
114
115  def next(self, value):
116    """
117    Provides the next enumeration after the given value.
118
119    :param str value: enumeration for which to get the next entry
120
121    :returns: enum value following the given entry
122
123    :raises: **ValueError** if no such element exists
124    """
125
126    if value not in self._values:
127      raise ValueError('No such enumeration exists: %s (options: %s)' % (value, ', '.join(self._values)))
128
129    next_index = (self._values.index(value) + 1) % len(self._values)
130    return self._values[next_index]
131
132  def previous(self, value):
133    """
134    Provides the previous enumeration before the given value.
135
136    :param str value: enumeration for which to get the previous entry
137
138    :returns: enum value proceeding the given entry
139
140    :raises: **ValueError** if no such element exists
141    """
142
143    if value not in self._values:
144      raise ValueError('No such enumeration exists: %s (options: %s)' % (value, ', '.join(self._values)))
145
146    prev_index = (self._values.index(value) - 1) % len(self._values)
147    return self._values[prev_index]
148
149  def __getitem__(self, item):
150    """
151    Provides the values for the given key.
152
153    :param str item: key to be looked up
154
155    :returns: **str** with the value for the given key
156
157    :raises: **ValueError** if the key doesn't exist
158    """
159
160    if item in vars(self):
161      return getattr(self, item)
162    else:
163      keys = ', '.join(self.keys())
164      raise ValueError("'%s' isn't among our enumeration keys, which includes: %s" % (item, keys))
165
166  def __iter__(self):
167    """
168    Provides an ordered listing of the enums in this set.
169    """
170
171    for entry in self._values:
172      yield entry
173