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